diff options
Diffstat (limited to 'SpatialSoftMax.lua')
| -rw-r--r-- | SpatialSoftMax.lua | 14 |
1 files changed, 11 insertions, 3 deletions
diff --git a/SpatialSoftMax.lua b/SpatialSoftMax.lua index ee93749..e245f24 100644 --- a/SpatialSoftMax.lua +++ b/SpatialSoftMax.lua @@ -15,6 +15,12 @@ function SpatialSoftMax:__init(fast) end function SpatialSoftMax: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 @@ -23,11 +29,14 @@ function SpatialSoftMax: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 SpatialSoftMax:updateOutput(input) - assert(input:dim() == 4 and input:isContiguous()); self:createIODescriptors(input) errcheck('cudnnSoftmaxForward', cudnn.handle[cutorch.getDevice()-1], @@ -38,8 +47,7 @@ function SpatialSoftMax:updateOutput(input) end function SpatialSoftMax: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('cudnnSoftmaxBackward', cudnn.handle[cutorch.getDevice()-1], |