FindEx implementetation + refactoring, take 3

1 files changed, 22 insertions, 22 deletions

diff --git a/TemporalConvolution.lua b/TemporalConvolution.lua
index 001d192..0654a39 100644
--- a/TemporalConvolution.lua
+++ b/TemporalConvolution.lua
@@ -6,6 +6,8 @@ local TemporalConvolution, parent =
 --it is recommended to pass padding parameter to this routine and use cudnn implicit padding facilities.
 --limitation is that padding will be equal on both sides.
 
+local Convolution = cudnn.SpatialConvolution
+
 function TemporalConvolution:__init(inputFrameSize, outputFrameSize,
                             kH, dH, padH)
     local delayedReset = self.reset
@@ -14,7 +16,7 @@ function TemporalConvolution:__init(inputFrameSize, outputFrameSize,
     local nOutputPlane = outputFrameSize
     self.inputFrameSize = inputFrameSize
     self.outputFrameSize = outputFrameSize
-    cudnn.SpatialConvolution.__init(self, nInputPlane, nOutputPlane, kW, kH, 1, dH,0,padH)
+    Convolution.__init(self, nInputPlane, nOutputPlane, kW, kH, 1, dH,0,padH)
     self.weight = self.weight:view(nOutputPlane,inputFrameSize*kH)
     self.gradWeight = self.gradWeight:view(outputFrameSize, inputFrameSize*kH)
 --self.dW and self.kW now have different meaning than in nn.TemporalConvolution, because
@@ -22,30 +24,23 @@ function TemporalConvolution:__init(inputFrameSize, outputFrameSize,
 end
 
 function TemporalConvolution:createIODescriptors(input)
-    local sizeChanged = false
-    if not self.iDesc or not self.oDesc or
-        input:size(1) ~= self.iSize[1] or input:size(2) ~= self.iSize[2]
-    or input:size(3) ~= self.iSize[3] or input:size(4) ~= self.iSize[4] then
-       sizeChanged = true
-    end
-    cudnn.SpatialConvolution.createIODescriptors(self,input)
-    if sizeChanged then
-       self.oSize = self.output:size()
-    end
+   return Convolution.createIODescriptors(self, input)
 end
 
 function TemporalConvolution:fastest(mode)
-    self = cudnn.SpatialConvolution.fastest(self,mode)
-    return self
+    return Convolution.fastest(self, mode)
 end
 
 function TemporalConvolution:setMode(fmode, bdmode, bwmode)
-    self = cudnn.SpatialConvolution.setMode(self,fmode, bdmode, bwmode)
-    return self
+   return Convolution.setMode(self, fmode, bdmode, bwmode)
+end
+
+function TemporalConvolution:resetMode()
+   return Convolution.resetMode(self)
 end
 
 function TemporalConvolution:resetWeightDescriptors()
-    cudnn.SpatialConvolution.resetWeightDescriptors(self)
+    return Convolution.resetWeightDescriptors(self)
 end
 
 local function inputview(input)
@@ -57,13 +52,14 @@ local function inputview(input)
 end
 
 function TemporalConvolution:updateOutput(input)
+   find.get():verifyWorkspaceSize(self)
    local _input = inputview(input)
    assert(_input:size(4) == self.inputFrameSize,'invalid input frame size')
    self.buffer = self.buffer or input.new()
    self._output = self._output or input.new()
    if self.output:storage() then self._output:set(self.output:storage()) else self._output = self.output end
    if self.buffer:storage() then self.output:set(self.buffer:storage(), 1, self.output:size()) else self.output = self.buffer end
-   cudnn.SpatialConvolution.updateOutput(self,_input)
+   Convolution.updateOutput(self, _input)
    self.buffer = self.output:view(self.oSize):transpose(2,3)
    self.output  = self._output:resize(self.buffer:size()):copy(self.buffer)
    -- self.output here is always 4D, use input dimensions to properly view output
@@ -90,9 +86,10 @@ end
 
 function TemporalConvolution:updateGradInput(input, gradOutput)
    if not self.gradInput then return end
+   find.get():verifyWorkspaceSize(self)
    local _gradOutput = transposeGradOutput(gradOutput,self.buffer)
    local _input = inputview(input)
-   self.gradInput = cudnn.SpatialConvolution.updateGradInput(self,_input, _gradOutput)
+   self.gradInput = Convolution.updateGradInput(self, _input, _gradOutput)
    if input:dim()==3 then
       self.gradInput = self.gradInput:view(self.iSize[1],self.iSize[3],self.iSize[4])
    else
@@ -102,11 +99,12 @@ function TemporalConvolution:updateGradInput(input, gradOutput)
 end
 
 function TemporalConvolution:accGradParameters(input,gradOutput,scale)
-   --2d (4d) view of input
-   local _input = inputview(input)
-   -- transpose gradOutput (it will likely be transposed twice, hopefully, no big deal
+   find.get():verifyWorkspaceSize(self)
+--2d (4d) view of input
+    local _input = inputview(input)
+-- transpose gradOutput (it will likely be transposed twice, hopefully, no big deal
     local _gradOutput = transposeGradOutput(gradOutput,self.buffer)
-    cudnn.SpatialConvolution.accGradParameters(self,_input,_gradOutput,scale)
+    return Convolution.accGradParameters(self,_input,_gradOutput,scale)
 end
 
 function TemporalConvolution:clearDesc()
@@ -130,3 +128,5 @@ function TemporalConvolution:clearState()
    nn.utils.clear(self, '_input', '_gradOutput')
    return parent.clearState(self)
 end
+
+return TemporalConvolution