CMul optimizations, doc and unit tests

3 files changed, 117 insertions, 33 deletions

diff --git a/CMul.lua b/CMul.lua
index 7cedfd2..ea3485a 100644
--- a/CMul.lua
+++ b/CMul.lua
@@ -1,14 +1,25 @@
 local CMul, parent = torch.class('nn.CMul', 'nn.Module')
 
-function CMul:__init(inputSize)
+function CMul:__init(...)
    parent.__init(self)
+   
+   local arg = {...}
+
+   self.size = torch.LongStorage()
+   local n = #arg
+   if n == 1 and torch.type(arg[1]) == 'torch.LongStorage' then
+      self.size:resize(#arg[1]):copy(arg[1])
+   else
+      self.size:resize(n)
+      for i=1,n do
+         self.size[i] = arg[i]
+      end
+   end
   
-   self.weight = torch.Tensor(inputSize)
-   self.gradWeight = torch.Tensor(inputSize)
+   self.weight = torch.Tensor(self.size)
+   self.gradWeight = torch.Tensor(self.size)
    
-   -- state
-   -- self.gradInput:resize(inputSize)
-   self.output:resize(inputSize) 
+   self.output:resize(self.size) 
 
    self:reset()
 end
@@ -17,52 +28,101 @@ function CMul:reset(stdv)
    if stdv then
       stdv = stdv * math.sqrt(3)
    else
-      stdv = 1./math.sqrt(self.weight:size(1))
+      stdv = 1./math.sqrt(self.weight:nElement())
    end
    self.weight:uniform(-stdv,stdv)
 end
 
 function CMul:updateOutput(input)
+   -- lazy-initialize
+   self._output = self._output or input.new()
+   self._weight = self._weight or input.new()
+   self._expand = self._expand or input.new()
+   self._repeat = self._repeat or input.new()
+   
    self.output:resizeAs(input):copy(input)
    if input:nElement() == self.weight:nElement() then
-      self.output:view(-1):cmul(self.weight:view(-1));
+      self._output:view(self.output, -1)
+      self._weight:view(self.weight, -1)
+      
+      self._output:cmul(self._weight)
    else
-      if input:isSameSizeAs(self.weight) then
-         self.output:cmul(self.weight)
+      local batchSize = input:size(1)
+      self._output:view(self.output, batchSize, -1)
+      self._weight:view(self.weight, 1, -1)
+      
+      self._expand:expandAs(self._weight, self._output)
+      
+      if torch.type(input) == 'torch.CudaTensor' then
+         self._repeat:resizeAs(self._expand):copy(self._expand)
+         self._output:cmul(self._repeat)
       else
-         local batchSize = input:size(1)
-         self.output:view(batchSize, -1):cmul(self.weight:view(1,-1):expandAs(input:view(batchSize, -1)))
+         self._output:cmul(self._expand)
       end
    end
+   
    return self.output
 end
 
 function CMul:updateGradInput(input, gradOutput)
-   if self.gradInput then
-      local nElement = self.gradInput:nElement()
-      self.gradInput:resizeAs(input)
-      self.gradInput:zero()
-      if self.weight:nElement() == gradOutput:nElement() then
-         self.gradInput:addcmul(1, self.weight, gradOutput)
+   if not self.gradInput then
+      return
+   end
+   
+   self._gradOutput = self._gradOutput or input.new()
+   self._gradInput = self._gradInput or input.new()
+
+   self.gradInput:resizeAs(input):zero()
+   if self.weight:nElement() == gradOutput:nElement() then
+      self.gradInput:addcmul(1, self.weight, gradOutput)
+   else
+      local batchSize = input:size(1)
+      self._gradOutput:view(gradOutput, batchSize, -1)
+      self._gradInput:view(self.gradInput, batchSize, -1)
+      self._weight:view(self.weight, 1, -1)
+      self._expand:expandAs(self._weight, self._gradOutput)
+      
+      if torch.type(input) == 'torch.CudaTensor' then
+         self._repeat:resizeAs(self._expand):copy(self._expand)
+         self._gradInput:addcmul(1, self._repeat, self._gradOutput)
       else
-         local gradOutput = gradOutput:view(input:size(1), -1)
-         local gradInput = self.gradInput:view(input:size(1), -1)
-         gradInput:addcmul(1, self.weight:view(1,-1):expandAs(gradOutput), gradOutput)
+         self._gradInput:addcmul(1, self._expand, self._gradOutput)
       end
-      return self.gradInput
    end
+   
+   return self.gradInput
 end
 
 function CMul:accGradParameters(input, gradOutput, scale)
+   scale = scale or 1
+   
+   self._input = self._input or input.new()
+   self._gradWeight = self._gradWeight or input.new()
+   self._sum = self._sum or input.new()
+   
    if self.weight:nElement() == gradOutput:nElement() then
-      self.gradWeight:addcmul(scale or 1, input, gradOutput)
+      self.gradWeight:addcmul(scale, input, gradOutput)
    else
       local batchSize = input:size(1)
-      local input = input:view(batchSize, -1)
-      local gradOutput = gradOutput:view(batchSize, -1)
-      local gradWeight = self.gradWeight:view(1, -1)
-      for i=1,batchSize do
-         gradWeight:addcmul(scale or 1, input[i], gradOutput[i])
-      end
+      self._input:view(input, batchSize, -1)
+      self._gradOutput:view(gradOutput, batchSize, -1)
+      self._gradWeight:view(self.gradWeight, 1, -1)
+      
+      self._repeat:cmul(self._input, self._gradOutput)
+      self._sum:sum(self._repeat, 1)
+      self._gradWeight:add(scale, self._sum)
+   end
+end
+
+function CMul:type(type)
+   if type then
+      self._input = nil
+      self._output = nil
+      self._weight = nil
+      self._gradWeight = nil
+      self._expand = nil
+      self._repeat = nil
+      self._sum = nil
    end
+   return parent.type(self, type)
 end
diff --git a/doc/simple.md b/doc/simple.md
index 4b1b182..e51ebbf 100644
--- a/doc/simple.md
+++ b/doc/simple.md
@@ -302,10 +302,12 @@ pi.
 <a name='nn.CMul'/>
 ## CMul ##
 
-`module` = `CMul(inputDimension)`
+`module` = `CMul(size)`
 
 Applies a component-wise multiplication to the incoming data, i.e.
-`y_i` = `w_i` =x_i=. 
+`y_i = w_i * x_i`. Argument `size` can be one or many numbers (sizes)
+or a `torch.LongStorage`. For example, `nn.CMul(3,4,5)` is equivalent to 
+`nn.CMul(torch.LongStorage{3,4,5})`.
 
 Example:
 ```lua
diff --git a/test.lua b/test.lua
index 828e218..27c3dde 100644
--- a/test.lua
+++ b/test.lua
@@ -74,7 +74,7 @@ function nntest.CMul()
    local inj = math.random(3,5)
    local ink = math.random(3,5)
    local input = torch.Tensor(ini,inj,ink):zero()
-   local module = nn.CMul(ini*inj*ink)
+   local module = nn.CMul(ini, inj, ink)
 
    -- 1D
    local err = jac.testJacobian(module,input)
@@ -94,7 +94,29 @@ function nntest.CMul()
    -- 2D
    local nframe = math.random(50,70)
    local nframe = 5
-   local input = torch.Tensor(nframe, ini,inj,ink):zero()
+   local input = torch.randn(nframe, ini,inj,ink)
+   local output = module:forward(input)
+   local output2 = torch.cmul(input, module.weight:view(1,ini,inj,ink):expandAs(input))
+   mytester:assertTensorEq(output2, output, 0.000001, 'CMul forward 2D err')
+   
+   module:zeroGradParameters()
+   local gradWeight = module.gradWeight:clone()
+   local gradInput = module:backward(input, output)
+   local gradInput2 = gradInput:clone():zero()
+   local outputView = output:view(input:size(1), -1)
+   gradInput2:view(input:size(1), -1):addcmul(1, module.weight:view(1,-1):expandAs(outputView), outputView)
+   mytester:assertTensorEq(gradInput2, gradInput, 0.000001, 'CMul updateGradInput 2D err')
+   mytester:assert(gradInput:isSameSizeAs(input), 'CMul gradInput 2D size err')
+   
+   local inputView = input:view(nframe, -1)
+   local gradWeightView = gradWeight:view(1, -1)
+   for i=1,nframe do
+      gradWeightView:addcmul(1, inputView[i], outputView[i])
+   end
+   mytester:assertTensorEq(gradWeight, module.gradWeight, 0.000001, 'CMul accGradParameters 2D err')
+   mytester:assert(module.weight:isSameSizeAs(module.gradWeight), 'CMul gradWeight size err')
+   
+   input:zero()
 
    local err = jac.testJacobian(module,input)
    mytester:assertlt(err,precision, 'error on state ')