WeightedEuclidean batch mode

3 files changed, 171 insertions, 39 deletions

diff --git a/WeightedEuclidean.lua b/WeightedEuclidean.lua
index 3808db6..5a3af27 100644
--- a/WeightedEuclidean.lua
+++ b/WeightedEuclidean.lua
@@ -6,13 +6,10 @@ function WeightedEuclidean:__init(inputSize,outputSize)
    self.templates = torch.Tensor(inputSize,outputSize)
    self.gradTemplates = torch.Tensor(inputSize,outputSize)
 
+   -- each template (output dim) has its own diagonal covariance matrix
    self.diagCov = torch.Tensor(inputSize,outputSize)
    self.gradDiagCov = torch.Tensor(inputSize,outputSize)
-
-   self.gradInput:resize(inputSize)
-   self.output:resize(outputSize)
-   self.temp = torch.Tensor(inputSize)
-
+   
    -- for compat with Torch's modules (it's bad we have to do that)
    do
       self.weight = self.templates
@@ -43,45 +40,137 @@ function WeightedEuclidean:reset(stdv)
 end
 
 function WeightedEuclidean:updateOutput(input)
-   self.output:zero()
-   for o = 1,self.templates:size(2) do
-      self.temp:copy(input):add(-1,self.templates:select(2,o))
-      self.temp:cmul(self.temp)
-      self.temp:cmul(self.diagCov:select(2,o)):cmul(self.diagCov:select(2,o))
-      self.output[o] = math.sqrt(self.temp:sum())
+   -- lazy-initialize 
+   self._temp = self._temp or self.output.new()
+   self._ones = self._ones or self.output.new{1}
+   self._diagCov = self._diagCov or self.output.new()
+   self._repeat = self._repeat or self.output.new()
+   self._sum = self._sum or self.output.new()
+   self._temp:resizeAs(input)
+   if input:dim() == 1 then
+      self.output:resize(self.templates:size(2))
+      for outIdx = 1,self.templates:size(2) do
+         self._temp:copy(input):add(-1,self.templates:select(2,outIdx))
+         self._temp:cmul(self._temp)
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._temp:cmul(diagCov):cmul(diagCov)
+         self.output[outIdx] = math.sqrt(self._temp:sum())
+      end
+   elseif input:dim() == 2 then
+      self.output:resize(input:size(1), self.templates:size(2))
+      if self._ones:size(1) ~= input:size(1) then
+         self._ones:resize(input:size(1)):fill(1)
+      end
+      for outIdx = 1,self.templates:size(2) do
+         self._temp:copy(input)
+         self._temp:addr(-1, self._ones, self.templates:select(2,outIdx))
+         self._temp:cmul(self._temp)
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._diagCov:resizeAs(diagCov):copy(diagCov)
+         self._diagCov:pow(2)
+         self._diagCov:resize(1,self._diagCov:size(1))
+         self._repeat:repeatTensor(self._diagCov, input:size(1), 1)
+         self._temp:cmul(self._temp, self._repeat)
+         self._sum:sum(self._temp, 2):sqrt()
+         self.output:select(2,outIdx):copy(self._sum)
+      end
+   else
+      error"1D or 2D input expected"
    end
    return self.output
 end
 
 function WeightedEuclidean:updateGradInput(input, gradOutput)
-   self:forward(input)
-   self.gradInput:zero()
-   for o = 1,self.templates:size(2) do
-      if self.output[o] ~= 0 then
-         self.temp:copy(input):add(-1,self.templates:select(2,o))
-         self.temp:cmul(self.diagCov:select(2,o)):cmul(self.diagCov:select(2,o))
-         self.temp:mul(gradOutput[o]/self.output[o])
-         self.gradInput:add(self.temp)
+   self._gradTemp = self._gradTemp or self.output.new()
+   self.gradInput:resizeAs(input):zero()
+   self._temp:resizeAs(input)
+   self._gradTemp:cdiv(gradOutput, self.output)
+   if input:dim() == 1 then
+      for outIdx = 1,self.templates:size(2) do
+         self._temp:copy(input):add(-1,self.templates:select(2,outIdx))
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._temp:cmul(diagCov):cmul(diagCov)
+         
+         self._temp:mul(self._gradTemp[outIdx])
+         self.gradInput:add(self._temp)
+      end
+   elseif input:dim() == 2 then
+      if self._ones:size(1) ~= input:size(1) then
+         self._ones:resize(input:size(1)):fill(1)
       end
+      for outIdx = 1,self.templates:size(2) do
+         self._temp:copy(input)
+         self._temp:addr(-1, self._ones, self.templates:select(2,outIdx))
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._diagCov:resizeAs(diagCov):copy(diagCov)
+         self._diagCov:pow(2)
+         self._diagCov:resize(1,self._diagCov:size(1))
+         self._repeat:repeatTensor(self._diagCov, input:size(1), 1)
+         self._temp:cmul(self._temp, self._repeat)
+
+         local gradTemp = self._gradTemp:select(2, outIdx)
+         gradTemp = gradTemp:reshape(1,gradTemp:size(1))
+         self._repeat:repeatTensor(gradTemp, input:size(2), 1)
+         self.gradInput:addcmul(1, self._temp, self._repeat)
+      end
+   else
+      error"1D or 2D input expected"
    end
    return self.gradInput
 end
 
 function WeightedEuclidean:accGradParameters(input, gradOutput, scale)
-   self:forward(input)
    scale = scale or 1
-   for o = 1,self.templates:size(2) do
-      if self.output[o] ~= 0 then
-         self.temp:copy(self.templates:select(2,o)):add(-1,input)
-         self.temp:cmul(self.diagCov:select(2,o)):cmul(self.diagCov:select(2,o))
-         self.temp:mul(gradOutput[o]/self.output[o])
-         self.gradTemplates:select(2,o):add(scale, self.temp)
+   self._temp:resizeAs(input)
+   self._gradTemp:cdiv(gradOutput, self.output)
+   if input:dim() == 1 then
+      for outIdx = 1,self.templates:size(2) do
+         self._temp:copy(self.templates:select(2,outIdx)):add(-1,input)
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._temp:cmul(diagCov):cmul(diagCov)
+         
+         self._temp:mul(self._gradTemp[outIdx])
+         self.gradTemplates:select(2,outIdx):add(scale, self._temp)
 
-         self.temp:copy(self.templates:select(2,o)):add(-1,input)
-         self.temp:cmul(self.temp)
-         self.temp:cmul(self.diagCov:select(2,o))
-         self.temp:mul(gradOutput[o]/self.output[o])
-         self.gradDiagCov:select(2,o):add(scale, self.temp)
+         self._temp:copy(self.templates:select(2,outIdx)):add(-1,input)
+         self._temp:pow(2)
+         self._temp:cmul(self.diagCov:select(2,outIdx))
+         self._temp:mul(self._gradTemp[outIdx])
+         self.gradDiagCov:select(2,outIdx):add(scale, self._temp)
       end
+   elseif input:dim() == 2 then
+      for outIdx = 1,self.templates:size(2) do
+         -- gradTemplates
+         self._temp:copy(input)
+         self._temp:addr(-1, self._ones, self.templates:select(2,outIdx))
+         local diagCov = self.diagCov:select(2,outIdx)
+         self._diagCov:resizeAs(diagCov):copy(diagCov)
+         self._diagCov:pow(2)
+         self._diagCov:resize(1,self._diagCov:size(1))
+         self._repeat:repeatTensor(self._diagCov, input:size(1), 1)
+         self._temp:cmul(self._temp, self._repeat)
+         
+         local gradTemp = self._gradTemp:select(2, outIdx)
+         gradTemp = gradTemp:reshape(1,gradTemp:size(1))
+         self._repeat:repeatTensor(gradTemp, input:size(2), 1)
+         self._temp:cmul(self._repeat)
+         self._sum:sum(self._temp, 1)
+         self.gradTemplates:select(2,outIdx):add(scale, self._sum)
+
+         -- gradDiagCov
+         local template = self.templates:select(2,outIdx)
+         template = template:reshape(1, template:size(1))
+         self._temp:repeatTensor(template, input:size(1), 1)
+         self._temp:add(-1,input)
+         self._temp:pow(2)
+         self._temp:cmul(self._repeat)
+         diagCov = diagCov:reshape(1, diagCov:size(1))
+         self._repeat:repeatTensor(self._diagCov, input:size(1), 1)
+         self._temp:cmul(self._repeat)
+         self._sum:sum(self._temp, 1)
+         self.gradDiagCov:select(2,outIdx):add(scale, self._sum)
+      end
+   else
+      error"1D or 2D input expected"
    end
 end
diff --git a/doc/simple.md b/doc/simple.md
index 2cc3ccf..8cbb017 100644
--- a/doc/simple.md
+++ b/doc/simple.md
@@ -387,22 +387,30 @@ then an `nxq` matrix would be output.
 <a name="nn.Euclidean"/>
 ## Euclidean ##
 
-`module` = `Euclidean(inputDimension,outputDimension)`
+`module` = `Euclidean(inputSize,outputSize)`
 
-Outputs the Euclidean distance of the input to `outputDimension` centers,
-i.e. this layer has the weights `c_i`, `i` = `1`,..,`outputDimension`, where
-`c_i` are vectors of dimension `inputDimension`. Output dimension `j` is
-`|| c_j - x ||`, where `x` is the input.
+Outputs the Euclidean distance of the input to `outputSize` centers,
+i.e. this layer has the weights `w_j`,  for `j` = `1`,..,`outputSize`, where
+`w_j` are vectors of dimension `inputSize`. 
+
+The distance `y_j` between center `j` and input `x` is formulated as
+`y_j = || w_j - x ||`.
 
 <a name="nn.WeightedEuclidean"/>
 ## WeightedEuclidean ##
 
-`module` = `WeightedEuclidean(inputDimension,outputDimension)`
+`module` = `WeightedEuclidean(inputSize,outputSize)`
 
 This module is similar to [Euclidean](#nn.Euclidean), but
 additionally learns a separate diagonal covariance matrix across the
-features of the input space for each center.
+features of the input space _for each center_. 
+
+In other words, for each of the `outputSize` centers `w_j`, there is 
+a diagonal covariance matrices `c_j`, for `j` = `1`,..,`outputSize`, 
+where `c_j` are stored as vectors of size `inputSize`.
 
+The distance `y_j` between center `j` and input `x` is formulated as
+`y_j = || c_j * (w_j - x) ||`.
 
 <a name="nn.Identity"/>
 ## Identity ##
diff --git a/test.lua b/test.lua
index 89ff7d5..c3c6cd3 100644
--- a/test.lua
+++ b/test.lua
@@ -456,7 +456,7 @@ function nntest.WeightedEuclidean()
    local inj = math.random(13,5)
    local input = torch.Tensor(ini):zero()
    local module = nn.WeightedEuclidean(ini,inj)
-
+   
    local err = jac.testJacobian(module,input)
    mytester:assertlt(err,precision, 'error on state ')
 
@@ -469,6 +469,41 @@ function nntest.WeightedEuclidean()
    local ferr,berr = jac.testIO(module,input)
    mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err ')
    mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err ')
+   
+   -- test batch
+   local bs = math.random(3,5)
+   input:uniform(0,1)
+   local output = module:forward(input):clone()
+   module:zeroGradParameters()
+   local gradInput = module:backward(input, output):clone()
+   local params, gradParams = module:parameters()
+   for i=1,#params do
+      params[i] = params[i]:clone()
+   end
+   local input2 = input:view(1, -1):repeatTensor(bs, 1)
+   local output2 = module:forward(input2)
+   module:zeroGradParameters()
+   local gradInput2 = module:backward(input2, output2, 1/bs)
+   local params2, gradParams2 = module:parameters()
+   mytester:assertTensorEq(output2[bs-1], output, 0.000001, "error in batch updateOutput")
+   mytester:assertTensorEq(gradInput2[bs-1], gradInput, 0.000001, "error in batch updateGradInput")
+   mytester:assertTensorEq(gradParams[1], gradParams2[1], 0.000001, "error in batch accGradParameters (gradTemplates)")
+   mytester:assertTensorEq(gradParams[2], gradParams2[2], 0.000001, "error in batch accGradParameters (gradDiagCov)")
+   
+   input:zero()
+   module:zeroGradParameters()
+   local err = jac.testJacobian(module,input)
+   mytester:assertlt(err,precision, 'error on state ')
+
+   local err = jac.testJacobianParameters(module, input, module.weight, module.gradWeight)
+   mytester:assertlt(err,precision, 'error on weight ')
+
+   local err = jac.testJacobianParameters(module, input, module.bias, module.gradBias)
+   mytester:assertlt(err,precision, 'error on bias ')
+
+   local ferr,berr = jac.testIO(module,input2)
+   mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err ')
+   mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err ')
 end
 
 local function criterionJacobianTest1D(cri, input, target)