diff options
Diffstat (limited to 'ReLU.lua')
-rw-r--r-- | ReLU.lua | 22 |
1 files changed, 15 insertions, 7 deletions
@@ -5,11 +5,17 @@ local errcheck = cudnn.errcheck function ReLU:__init() parent.__init(self) - self.iSize = torch.LongStorage(4):fill(0) + self.iSize = torch.LongStorage(4):fill(0) end function ReLU:createIODescriptors(input) - if not self.iDesc or not self.oDesc or + 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() @@ -17,26 +23,28 @@ function ReLU:createIODescriptors(input) self.output:resizeAs(input) 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 ReLU:updateOutput(input) - assert(input:dim() == 4 and input:isContiguous()); self:createIODescriptors(input) errcheck('cudnnActivationForward', cudnn.handle[cutorch.getDevice()-1], 'CUDNN_ACTIVATION_RELU', - self.iDesc[0], input:data(), + self.iDesc[0], input:data(), self.oDesc[0], self.output:data()); return self.output end function ReLU:updateGradInput(input, gradOutput) - assert(input:dim() == 4 and input:isContiguous()); - assert(gradOutput:dim() == 4 and gradOutput:isContiguous()); + assert((gradOutput:dim() == 4 or gradOutput:dim() == 3) and gradOutput:isContiguous()); self:createIODescriptors(input) errcheck('cudnnActivationBackward', cudnn.handle[cutorch.getDevice()-1], 'CUDNN_ACTIVATION_RELU', self.oDesc[0], self.output:data(), self.oDesc[0], gradOutput:data(), - self.iDesc[0], input:data(), + self.iDesc[0], input:data(), self.iDesc[0], self.gradInput:data()); return self.gradInput end |