diff options
author | soumith <soumith@fb.com> | 2014-11-26 05:23:52 +0300 |
---|---|---|
committer | soumith <soumith@fb.com> | 2014-11-26 05:23:52 +0300 |
commit | b0e6e3f0ad10e931a1f83197b127a442e179e67e (patch) | |
tree | 4d58247a109e9a211317625e1263a637a5981070 /SpatialMaxPooling.lua | |
parent | b38b1a17251639eb68df4acde5784a661e2c8888 (diff) |
refactoring lots of duplicated code
Diffstat (limited to 'SpatialMaxPooling.lua')
-rw-r--r-- | SpatialMaxPooling.lua | 78 |
1 files changed, 2 insertions, 76 deletions
diff --git a/SpatialMaxPooling.lua b/SpatialMaxPooling.lua index 8bd1bf6..108a055 100644 --- a/SpatialMaxPooling.lua +++ b/SpatialMaxPooling.lua @@ -1,80 +1,6 @@ -local SpatialMaxPooling, parent = torch.class('cudnn.SpatialMaxPooling', 'nn.Module') -local ffi = require 'ffi' -local C = cudnn.C -local errcheck = cudnn.errcheck +local SpatialMaxPooling, parent = torch.class('cudnn.SpatialMaxPooling', 'cudnn._Pooling') function SpatialMaxPooling:__init(kW, kH, dW, dH) - parent.__init(self) - self.kW = kW - self.kH = kH - self.dW = dW or kW - self.dH = dH or kW - self.iSize = torch.LongStorage(4):fill(0) + parent.__init(self, kW, kH, dW, dH) self.mode = 'CUDNN_POOLING_MAX' end - -function SpatialMaxPooling:resetPoolDescriptors() - -- create pooling descriptor - self.poolDesc = ffi.new('struct cudnnPoolingStruct*[1]') - errcheck('cudnnCreatePoolingDescriptor', self.poolDesc) - errcheck('cudnnSetPoolingDescriptor', self.poolDesc[0], self.mode, - self.kH, self.kW, self.dH, self.dW); - local function destroyPoolDesc(d) - errcheck('cudnnDestroyPoolingDescriptor', d[0]); - end - ffi.gc(self.poolDesc, destroyPoolDesc) -end - -function SpatialMaxPooling:createIODescriptors(input) - local batch = true - if input:dim() == 3 then - input = input:view(1, input:size(1), input:size(2), input:size(3)) - batch = false - end - assert(input:dim() == 4 and input:isContiguous()); - 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 - self.iSize = input:size() - -- resize gradInput - self.gradInput:resizeAs(input) - -- resize output - local oW = math.floor((input:size(4) - self.kW)/self.dW + 1) - local oH = math.floor((input:size(3) - self.kH)/self.dH + 1) - self.output:resize(input:size(1), input:size(2), oH, oW) - - -- create input/output descriptor - self.iDesc = cudnn.toDescriptor(input) - self.oDesc = cudnn.toDescriptor(self.output) - if not batch then - self.gradInput = self.gradInput:view(self.gradInput:size(2), self.gradInput:size(3), self.gradInput:size(4)) - self.output = self.output:view(self.output:size(2), self.output:size(3), self.output:size(4)) - end - end -end - -function SpatialMaxPooling:updateOutput(input) - if not self.poolDesc then self:resetPoolDescriptors() end - self:createIODescriptors(input) - errcheck('cudnnPoolingForward', cudnn.handle[cutorch.getDevice()-1], self.poolDesc[0], - self.iDesc[0], input:data(), - self.oDesc[0], self.output:data()); - return self.output -end - -function SpatialMaxPooling:updateGradInput(input, gradOutput) - assert(gradOutput:dim() == 3 or gradOutput:dim() == 4); - if not gradOutput:isContiguous() then - self._gradOutput = self._gradOutput or gradOutput.new() - self._gradOutput:resizeAs(gradOutput):copy(gradOutput) - gradOutput = self._gradOutput - end - if not self.poolDesc then self:resetPoolDescriptors() end - self:createIODescriptors(input) - errcheck('cudnnPoolingBackward', cudnn.handle[cutorch.getDevice()-1], self.poolDesc[0], - self.oDesc[0], self.output:data(), - self.oDesc[0], gradOutput:data(), - self.iDesc[0], input:data(), - self.iDesc[0], self.gradInput:data()); - return self.gradInput -end |