From 56b6d5426509b4d0bef7d2648fad72ab4c122c84 Mon Sep 17 00:00:00 2001
From: soumith <soumith@fb.com>
Date: Mon, 17 Nov 2014 18:59:09 -0800
Subject: adding non-batch mode

---
 SpatialMaxPooling.lua | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

(limited to 'SpatialMaxPooling.lua')

diff --git a/SpatialMaxPooling.lua b/SpatialMaxPooling.lua
index 53c93a9..8bd1bf6 100644
--- a/SpatialMaxPooling.lua
+++ b/SpatialMaxPooling.lua
@@ -19,14 +19,20 @@ function SpatialMaxPooling:resetPoolDescriptors()
    errcheck('cudnnCreatePoolingDescriptor', self.poolDesc)
    errcheck('cudnnSetPoolingDescriptor', self.poolDesc[0], self.mode,
             self.kH, self.kW, self.dH, self.dW);
-   local function destroyPoolDesc(d) 
+   local function destroyPoolDesc(d)
       errcheck('cudnnDestroyPoolingDescriptor', d[0]);
    end
    ffi.gc(self.poolDesc, destroyPoolDesc)
 end
 
 function SpatialMaxPooling: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()
@@ -40,22 +46,24 @@ function SpatialMaxPooling:createIODescriptors(input)
       -- 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)
-   assert(input:dim() == 4 and input:isContiguous());
    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.iDesc[0], input:data(),
             self.oDesc[0], self.output:data());
    return self.output
 end
 
 function SpatialMaxPooling:updateGradInput(input, gradOutput)
-   assert(input:dim() == 4 and input:isContiguous());
-   assert(gradOutput:dim() == 4);
+   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)
@@ -66,8 +74,7 @@ function SpatialMaxPooling:updateGradInput(input, gradOutput)
    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], input:data(),
             self.iDesc[0], self.gradInput:data());
    return self.gradInput
 end
-
-- 
cgit v1.2.3