Added new Optimization framework to support new algos.

5 files changed, 102 insertions, 15 deletions

diff --git a/Optimization.lua b/Optimization.lua
new file mode 100644
index 0000000..ed230e5
--- /dev/null
+++ b/Optimization.lua
@@ -0,0 +1,55 @@
+local Optimization = torch.class('nn.Optimization')
+
+function Optimization:__init()
+end
+
+function Optimization:forward(parameters, gradParameters)
+   self:flatten(parameters, gradParameters)
+   -- do your thing
+   self:unflatten(parameters, gradParameters)
+end
+
+function Optimization:flatten(parameters, gradParameters)
+   if type(parameters) == 'table' then
+      -- create flat parameters
+      self.parameters = self.parameters or torch.Tensor()
+      self.gradParameters = self.gradParameters or torch.Tensor()
+      -- assuming that the parameters won't change their size, 
+      -- we compute offsets once
+      if not self.offsets then
+         self.nParameters = 0
+         self.offsets = {}
+         for _,param in ipairs(parameters) do
+            table.insert(self.offsets, self.nParameters+1)
+            self.nParameters = self.nParameters + param:nElement()
+         end
+         self.parameters:resize(self.nParameters)
+         self.gradParameters:resize(self.nParameters)
+      end
+      -- copy all params in flat array
+      for i = 1,#parameters do
+         local nElement = parameters[i]:nElement()
+         self.parameters:narrow(1,self.offsets[i],nElement):copy(parameters[i])
+         self.gradParameters:narrow(1,self.offsets[i],nElement):copy(gradParameters[i])
+      end
+   else
+      self.parameters = parameters
+      self.gradParameters = gradParameters
+   end
+end
+
+function Optimization:unflatten(parameters, gradParameters)
+   if type(parameters) == 'table' then
+      -- copy all params into unflat arrays
+      local offset = 1
+      for i = 1,#parameters do
+         local nElement = parameters[i]:nElement()
+         parameters[i]:copy(self.parameters:narrow(1,offset,nElement))
+         gradParameters[i]:copy(self.gradParameters:narrow(1,offset,nElement))
+         offset = offset + nElement
+      end
+   else
+      parameters = self.parameters
+      gradParameters = self.gradParameters
+   end
+end
diff --git a/SGDOptimization.lua b/SGDOptimization.lua
new file mode 100644
index 0000000..e2dd875
--- /dev/null
+++ b/SGDOptimization.lua
@@ -0,0 +1,36 @@
+local SGD,parent = torch.class('nn.SGDOptimization', 'nn.Optimization')
+
+function SGD:__init(...)
+   parent.__init(self)
+   xlua.unpack_class(self, {...},
+      'SGDOptimization', nil,
+      {arg='learningRate', type='number', help='learning rate (W = W - rate*dE/dW)', default=1e-2},
+      {arg='weightDecay', type='number', help='amount of weight decay (W = W - decay*W)', default=0},
+      {arg='momentum', type='number', help='amount of momentum on weights (dE/W = dE/dW + momentum*prev(dE/dW))', default=0}
+   )
+end
+
+function SGD:forward(parameters, gradParameters)
+   self:flatten(parameters, gradParameters)
+
+   -- apply momentum
+   if self.momentum ~= 0 then
+      if not self.currentGradParameters then
+         self.currentGradParameters = torch.Tensor():resizeAs(self.gradParameters):copy(self.gradParameters)
+      else
+         self.currentGradParameters:mul(self.momentum):add(self.gradParameters):div(1+self.momentum)
+      end
+   else
+      self.currentGradParameters = self.gradParameters
+   end
+
+   -- weight decay
+   if self.weightDecay ~= 0 then
+      self.parameters:add(-self.weightDecay, self.parameters)
+   end
+
+   -- update parameters
+   self.parameters:add(-self.learningRate, self.currentGradParameters)
+
+   self:unflatten(parameters, gradParameters)
+end
diff --git a/StochasticTrainer.lua b/StochasticTrainer.lua
index 4d72acb..62fb670 100644
--- a/StochasticTrainer.lua
+++ b/StochasticTrainer.lua
@@ -32,6 +32,8 @@ function StochasticTrainer:__init(...)
       {arg='save', type='string', help='path to save networks and log training'},
       {arg='timestamp', type='boolean', help='if true, appends a timestamp to each network saved', default=false}
    )
+   -- instantiate SGD optimization module
+   self.optimizer = nn.SGDOptimization(self.learningRate, self.weightDecay, self.momentum)
    -- private params
    self.errorArray = self.skipUniformTargets
    self.trainOffset = 0
@@ -129,13 +131,7 @@ function StochasticTrainer:train(dataset)
             self.currentError = self.currentError + error
 
             -- reset gradients
-            if self.momentum ~= 0 then
-               for _,grad in ipairs(gradParameters) do
-                  grad:mul(self.momentum)
-               end
-            else
-               module:zeroGradParameters()
-            end
+            module:zeroGradParameters()
 
             -- backward through model
             -- (if no criterion, it is assumed that derror is internally generated)
@@ -146,15 +142,8 @@ function StochasticTrainer:train(dataset)
                module:backward(input)
             end
 
-            -- weight decay ?
-            if self.weightDecay ~= 0 then
-               for _,param in ipairs(parameters) do
-                  param:add(-self.weightDecay, param)
-               end
-            end
-
             -- update parameters in the model
-            module:updateParameters(currentLearningRate)
+            self.optimizer:forward(parameters, gradParameters)
          end
 
          -- call user hook, if any
@@ -178,6 +167,7 @@ function StochasticTrainer:train(dataset)
 
       self.epoch = self.epoch + 1
       currentLearningRate = self.learningRate/(1+self.epoch*self.learningRateDecay)
+      self.optimizer.learningRate = currentLearningRate
 
       if dataset.infiniteSet then
          self.trainOffset = self.trainOffset + dataset:size()
diff --git a/init.lua b/init.lua
index 99e2fa7..3fc2ac1 100644
--- a/init.lua
+++ b/init.lua
@@ -97,6 +97,10 @@ torch.include('nnx', 'SpatialMSECriterion.lua')
 torch.include('nnx', 'SpatialClassNLLCriterion.lua')
 torch.include('nnx', 'SpatialSparseCriterion.lua')
 
+-- optimizations:
+torch.include('nnx', 'Optimization.lua')
+torch.include('nnx', 'SGDOptimization.lua')
+
 -- trainers:
 torch.include('nnx', 'Trainer.lua')
 torch.include('nnx', 'StochasticTrainer.lua')
diff --git a/nnx-1.0-1.rockspec b/nnx-1.0-1.rockspec
index ed70f18..ba2120c 100644
--- a/nnx-1.0-1.rockspec
+++ b/nnx-1.0-1.rockspec
@@ -94,6 +94,8 @@ build = {
          install_files(/lua/nnx SpatialGraph.lua)
          install_files(/lua/nnx SpatialColorTransform.lua)
          install_files(/lua/nnx SpatialRecursiveFovea.lua)
+         install_files(/lua/nnx Optimization.lua)
+         install_files(/lua/nnx SGDOptimization.lua)
          add_subdirectory (test)
          install_targets(/lib nnx)
    ]],