Euclidean

2 files changed, 191 insertions, 23 deletions

diff --git a/Euclidean.lua b/Euclidean.lua
index 229d792..d5700a5 100644
--- a/Euclidean.lua
+++ b/Euclidean.lua
@@ -9,7 +9,8 @@ function Euclidean:__init(inputSize,outputSize)
    -- state
    self.gradInput:resize(inputSize)
    self.output:resize(outputSize)
-   self.temp = torch.Tensor(inputSize)
+   
+   self.fastBackward = true
 
    self:reset()
 end
@@ -31,37 +32,159 @@ function Euclidean:reset(stdv)
    end
 end
 
+local function view(res, src, ...)
+   local args = {...}
+   if src:isContiguous() then
+      res:view(src, unpack(args))
+   else
+      res:reshape(src, unpack(args))
+   end
+end
+
 function Euclidean:updateOutput(input)
-   self.output:zero()
-   for o = 1,self.weight:size(2) do
-      self.output[o] = input:dist(self.weight:select(2,o))
+   -- lazy initialize buffers
+   self._input = self._input or input.new()
+   self._weight = self._weight or self.weight.new()
+   self._expand = self._expand or self.output.new()
+   self._expand2 = self._expand2 or self.output.new()
+   self._repeat = self._repeat or self.output.new()
+   self._repeat2 = self._repeat2 or self.output.new()
+   
+   local inputSize, outputSize = self.weight:size(1), self.weight:size(2)
+   
+   -- y_j = || w_j - x || = || x - w_j ||
+   if input:dim() == 1 then
+      view(self._input, input, inputSize, 1)
+      self._expand:expandAs(self._input, self.weight)
+      self._repeat:resizeAs(self._expand):copy(self._expand)
+      self._repeat:add(-1, self.weight)
+      self.output:norm(self._repeat, 2, 1)
+      self.output:resize(outputSize)
+   elseif input:dim() == 2 then
+      local batchSize = input:size(1)
+      
+      view(self._input, input, batchSize, inputSize, 1)
+      self._expand:expand(self._input, batchSize, inputSize, outputSize)
+      -- make the expanded tensor contiguous (requires lots of memory)
+      self._repeat:resizeAs(self._expand):copy(self._expand)
+      
+      self._weight:view(self.weight, 1, inputSize, outputSize)
+      self._expand2:expandAs(self._weight, self._repeat)
+      
+      if torch.type(input) == 'torch.CudaTensor' then
+         -- requires lots of memory, but minimizes cudaMallocs and loops
+         self._repeat2:resizeAs(self._expand2):copy(self._expand2)
+         self._repeat:add(-1, self._repeat2)
+      else
+         self._repeat:add(-1, self._expand2)
+      end
+      
+      self.output:norm(self._repeat, 2, 2)
+      self.output:resize(batchSize, outputSize)
+   else
+      error"1D or 2D input expected"
    end
+   
    return self.output
 end
 
 function Euclidean:updateGradInput(input, gradOutput)
-   self:updateOutput(input)
-   if self.gradInput then
-      self.gradInput:zero()
-      for o = 1,self.weight:size(2) do
-         if self.output[o] ~= 0 then
-            self.temp:copy(input):add(-1,self.weight:select(2,o))
-            self.temp:mul(gradOutput[o]/self.output[o])
-            self.gradInput:add(self.temp)
-         end
+   if not self.gradInput then
+      return
+   end
+   
+   self._div = self._div or input.new()
+   self._output = self._output or self.output.new()
+   self._gradOutput = self._gradOutput or input.new()
+   self._expand3 = self._expand3 or input.new()
+   
+   if not self.fastBackward then
+      self:updateOutput(input)
+   end
+   
+   local inputSize, outputSize = self.weight:size(1), self.weight:size(2)
+   
+   --[[ 
+   dy_j   -2 * (w_j - x)     x - w_j
+   ---- = ---------------  = -------
+    dx    2 || w_j - x ||      y_j
+   --]]
+
+   -- to prevent div by zero (NaN) bugs
+   self._output:resizeAs(self.output):copy(self.output):add(0.0000001)
+   view(self._gradOutput, gradOutput, gradOutput:size())
+   self._div:cdiv(gradOutput, self._output) 
+   if input:dim() == 1 then
+      self._div:resize(1, outputSize)
+      self._expand3:expandAs(self._div, self.weight)
+      
+      if torch.type(input) == 'torch.CudaTensor' then
+         self._repeat2:resizeAs(self._expand3):copy(self._expand3)
+         self._repeat2:cmul(self._repeat)
+      else
+         self._repeat2:cmul(self._repeat, self._expand3)
+      end
+      
+      self.gradInput:sum(self._repeat2, 2)
+      self.gradInput:resizeAs(input)
+   elseif input:dim() == 2 then
+      local batchSize = input:size(1)
+      
+      self._div:resize(batchSize, 1, outputSize)
+      self._expand3:expand(self._div, batchSize, inputSize, outputSize)
+      
+      if torch.type(input) == 'torch.CudaTensor' then
+         self._repeat2:resizeAs(self._expand3):copy(self._expand3)
+         self._repeat2:cmul(self._repeat)
+      else
+         self._repeat2:cmul(self._repeat, self._expand3)
       end
-      return self.gradInput
+      
+      self.gradInput:sum(self._repeat2, 3)
+      self.gradInput:resizeAs(input)
+   else
+      error"1D or 2D input expected"
    end
+   
+   return self.gradInput
 end
 
 function Euclidean:accGradParameters(input, gradOutput, scale)
-   self:updateOutput(input)
+   local inputSize, outputSize = self.weight:size(1), self.weight:size(2)
    scale = scale or 1
-   for o = 1,self.weight:size(2) do
-      if self.output[o] ~= 0 then
-         self.temp:copy(self.weight:select(2,o)):add(-1,input)
-         self.temp:mul(gradOutput[o]/self.output[o])
-         self.gradWeight:select(2,o):add(scale, self.temp)
-      end
+   
+   --[[ 
+   dy_j    2 * (w_j - x)     w_j - x
+   ---- = ---------------  = -------
+   dw_j   2 || w_j - x ||      y_j
+   --]]
+   -- assumes a preceding call to updateGradInput 
+   if input:dim() == 1 then
+      self.gradWeight:add(-scale, self._repeat2)
+   elseif input:dim() == 2 then
+      self._sum = self._sum or input.new()
+      self._sum:sum(self._repeat2, 1)
+      self._sum:resize(inputSize, outputSize)
+      self.gradWeight:add(-scale, self._sum)
+   else
+      error"1D or 2D input expected"
+   end
+end
+
+function Euclidean:type(type)
+   if type then
+      -- prevent premature memory allocations
+      self._input = nil
+      self._output = nil
+      self._gradOutput = nil
+      self._weight = nil
+      self._div = nil
+      self._sum = nil
+      self._expand = nil
+      self._expand2 = nil
+      self._expand3 = nil
+      self._repeat = nil
+      self._repeat2 = nil
    end
+   return parent.type(self, type)
 end
diff --git a/test.lua b/test.lua
index 27c3dde..0005761 100644
--- a/test.lua
+++ b/test.lua
@@ -480,9 +480,54 @@ end
 function nntest.Euclidean()
    local ini = math.random(5,7)
    local inj = math.random(5,7)
-   local input = torch.Tensor(ini):zero()
+   local input = torch.randn(ini)
+   local gradOutput = torch.randn(inj)
    local module = nn.Euclidean(ini,inj)
-
+   local output = module:forward(input):clone()
+   
+   local output2 = torch.Tensor(inj):zero()
+   for o = 1,module.weight:size(2) do
+      output2[o] = input:dist(module.weight:select(2,o))
+   end
+   mytester:assertTensorEq(output, output2, 0.000001, 'Euclidean forward 1D err')
+   
+   local input2 = torch.randn(8, ini)
+   input2[2]:copy(input)
+   local output2 = module:forward(input2)
+   mytester:assertTensorEq(output2[2], output, 0.000001, 'Euclidean forward 2D err')
+   
+   local output = module:forward(input):clone()
+   module:zeroGradParameters()
+   local gradInput = module:backward(input, gradOutput, 1):clone()
+   local gradInput2 = torch.zeros(ini)
+   local temp = input:clone()
+   for o = 1,module.weight:size(2) do
+      temp:copy(input)
+      temp:add(-1,module.weight:select(2,o))
+      temp:mul(gradOutput[o]/output[o])
+      gradInput2:add(temp)
+   end
+   mytester:assertTensorEq(gradInput, gradInput2, 0.000001, 'Euclidean updateGradInput 1D err')
+   
+   local gradWeight = module.gradWeight:clone():zero()
+   for o = 1,module.weight:size(2) do
+      temp:copy(module.weight:select(2,o)):add(-1,input)
+      temp:mul(gradOutput[o]/output[o])
+      gradWeight:select(2,o):add(1, temp)
+   end
+   mytester:assertTensorEq(gradWeight, module.gradWeight, 0.000001, 'Euclidean accGradParameters 1D err')
+   
+   local input2 = input:view(1, -1):repeatTensor(8, 1)
+   local gradOutput2 = gradOutput:view(1, -1):repeatTensor(8, 1)
+   local output2 = module:forward(input2)
+   module:zeroGradParameters()
+   local gradInput2 = module:backward(input2, gradOutput2, 1/8)
+   mytester:assertTensorEq(gradInput2[2], gradInput, 0.000001, 'Euclidean updateGradInput 2D err')
+   
+   mytester:assertTensorEq(gradWeight, module.gradWeight, 0.000001, 'Euclidean accGradParameters 2D err')
+   
+   input:zero()
+   module.fastBackward = false
    local err = jac.testJacobian(module,input)
    mytester:assertlt(err,precision, 'error on state ')