use cudnn for temporal convolution

1 files changed, 103 insertions, 0 deletions

diff --git a/TemporalConvolution.lua b/TemporalConvolution.lua
new file mode 100644
index 0000000..9ace2f2
--- /dev/null
+++ b/TemporalConvolution.lua
@@ -0,0 +1,103 @@
+local TemporalConvolution, parent =
+    torch.class('cudnn.TemporalConvolution', 'cudnn.SpatialConvolution')
+--use cudnn to perform temporal convolutions
+--note: if padH parameter is not passed, no padding will be performed, as in parent TemporalConvolution
+--however, instead of separately padding data, as is required now for nn.TemporalConvolution, 
+--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.
+
+function TemporalConvolution:__init(inputFrameSize, outputFrameSize,
+                            kH, dH, padH)
+    local delayedReset = self.reset
+    local kW = inputFrameSize
+    local nInputPlane = 1 -- single channel
+    local nOutputPlane = outputFrameSize
+    parent.__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.inputFrameSize = inputFrameSize
+    self.outputFrameSize = outputFramesize
+--self.dW and self.kW now have different meaning than in nn.TemporalConvolution, because
+--W and H are switched in temporal and spatial
+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
+    parent.createIODescriptors(self,input)
+    if sizeChanged then
+       self.oSize = self.output:size()
+    end
+end
+
+local function inputview(input)
+   local _input = input
+   if input:dim()==2 then
+      _input = input:view(1,input:size(1),input:size(2))
+   end
+   return _input:view(_input:size(1),1,_input:size(2),_input:size(3))
+end
+
+function TemporalConvolution:updateOutput(input)
+   local _input = inputview(input)
+   assert(_input:size(4) == self.inputFrameSize,'invalid input frame size')  
+   self.buffer = self.buffer or torch.CudaTensor()
+   self._output = self._output or torch.CudaTensor()
+   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()) else self.output = self.buffer end
+   parent.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
+   if input:dim()==3 then
+     self.output=self.output:view(self.oSize[1], self.oSize[3],self.oSize[2])
+   else 
+     self.output=self.output:view(self.oSize[3], self.oSize[2])
+   end  
+   return self.output
+end
+
+local function transposeGradOutput(src,dst)
+    assert(src:dim() == 2 or src:dim() == 3, 'gradOutput has to be 2D or 3D');
+    local srctransposed = src:transpose(src:dim(),src:dim()-1)    
+    dst:resize(srctransposed:size())
+    dst:copy(srctransposed)
+    if src:dim()==3 then
+      dst = dst:view(dst:size(1),dst:size(2),dst:size(3),1)
+    else
+      dst = dst:view(dst:size(1),dst:size(2),1)
+    end      
+    return dst 
+end
+
+function TemporalConvolution:updateGradInput(input, gradOutput)
+   if not self.gradInput then return end
+   local _gradOutput = transposeGradOutput(gradOutput,self.buffer)
+   local _input = inputview(input)
+   self.gradInput = parent.updateGradInput(self,_input, _gradOutput)
+   if input:dim()==3 then
+      self.gradInput = self.gradInput:view(self.gradInput:size(1),self.gradInput:size(3),self.gradInput:size(4))
+   else
+      self.gradInput = self.gradInput:view(self.gradInput:size(3),self.gradInput:size(4))
+   end
+   return self.gradInput
+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
+    local _gradOutput = transposeGradOutput(gradOutput,self.buffer)
+    parent.accGradParameters(self,_input,_gradOutput,scale)
+end
+
+function TemporalConvolution:write(f)
+  self.buffer = nil
+  self._ouptut = nil
+  self.oSize = nil
+  parent.write(self,f)
+end