added allreduce to lBFGS

3 files changed, 23 insertions, 14 deletions

diff --git a/BatchOptimization.lua b/BatchOptimization.lua
index 1e83bbb..cbc55e3 100644
--- a/BatchOptimization.lua
+++ b/BatchOptimization.lua
@@ -17,7 +17,13 @@ function Batch:__init(...)
                      {arg='precode', type='function',
                       help='optional code to be run by each parallel worker at their init'},
                      {arg='verbose', type='number',
-                      help='verbose level during training [0-2]', default=0}
+                      help='verbose level during training [0-2]', default=0},
+                     {arg='allreduce', type='boolean', help='use allreduce', default=false},
+                     {arg='allreduceSyncTime', type='boolean', help='sync period', default=1},
+                     {arg='allreduceMaster', type='string', help='master address', default='localhost'},
+                     {arg='allreduceUniqueId', type='boolean', help='job unique id', default=0},
+                     {arg='allreduceNbNodes', type='boolean', help='number of nodes', default=1},
+                     {arg='allreduceNodeId', type='boolean', help='this node\'s id', default=1}
                   )
    self.parameters = nnx.flattenParameters(nnx.getParameters(self.module))
    self.gradParameters = nnx.flattenParameters(nnx.getGradParameters(self.module))
@@ -30,6 +36,13 @@ function Batch:__init(...)
       self:setup_mapreduce()
    end
    self.P = self.parallelize
+
+   if self.allreduce then
+      require 'allreduce'
+      allreduce.init(self.allreduceMaster, self.allreduceUniqueId, 
+                     self.allreduceNbNodes, self.allreduceNodeId)
+      self.accError = 0
+   end
 end
 
 function Batch:forward(inputs, targets, options)
diff --git a/LBFGSOptimization.lua b/LBFGSOptimization.lua
index 53dfe70..bf9d41e 100644
--- a/LBFGSOptimization.lua
+++ b/LBFGSOptimization.lua
@@ -38,6 +38,14 @@ end
 function LBFGS:optimize()
    -- callback for lBFGS
    lbfgs.evaluate = self.evaluate
+
+   -- allreduce sync
+   if self.allreduce then
+      if (self.sampleCounter % self.allreduceSyncTime) == self.allreduceSyncTime-1 then
+         allreduce.accumulate(self.parameters)
+      end
+   end
+
    -- the magic function: will update the parameter vector according to the l-BFGS method
    self.output = lbfgs.run()
 end
diff --git a/SGDOptimization.lua b/SGDOptimization.lua
index 09b9243..1309ae9 100644
--- a/SGDOptimization.lua
+++ b/SGDOptimization.lua
@@ -13,20 +13,8 @@ function SGD:__init(...)
       {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*(1-momentum) + prev(dE/dW)*momentum)', default=0},
-      {arg='allreduce', type='boolean', help='use allreduce', default=false},
-      {arg='allreduceSyncTime', type='boolean', help='sync period', default=10},
-      {arg='allreduceMaster', type='string', help='master address', default='localhost'},
-      {arg='allreduceUniqueId', type='boolean', help='job unique id', default=0},
-      {arg='allreduceNbNodes', type='boolean', help='number of nodes', default=1},
-      {arg='allreduceNodeId', type='boolean', help='this node\'s id', default=1}
+       help='amount of momentum on weights (dE/W = dE/dW*(1-momentum) + prev(dE/dW)*momentum)', default=0}
    )
-   if self.allreduce then
-      require 'allreduce'
-      allreduce.init(self.allreduceMaster, self.allreduceUniqueId, 
-                     self.allreduceNbNodes, self.allreduceNodeId)
-      self.accError = 0
-   end
 end
 
 function SGD:optimize()