make avg-pooling swappable with nn

1 files changed, 25 insertions, 3 deletions

diff --git a/SpatialAveragePooling.lua b/SpatialAveragePooling.lua
index e94affe..d25138d 100644
--- a/SpatialAveragePooling.lua
+++ b/SpatialAveragePooling.lua
@@ -1,9 +1,31 @@
 local SpatialAveragePooling, parent
    = torch.class('cudnn.SpatialAveragePooling', 'cudnn._Pooling')
 
-function SpatialAveragePooling:__init(kW, kH, dW, dH, padW, padH)
-   parent.__init(self, kW, kH, dW, dH, padW, padH)
-   self.mode = 'CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING'
+local function backwardCompatible(self)
+   if self.ceil_mode == nil then
+      self.ceil_mode = false
+      self.count_include_pad = true
+      self.padH = 0
+      self.padW = 0
+   end
+end
+
+function SpatialAveragePooling:updateOutput(input)
+   -- for nn <> cudnn conversion
+   backwardCompatible(self)
+   if self.divide ~= nil then
+      assert(self.divide, 'not supported')
+   end
+
+   self.count_include_pad = self.count_include_pad ~= nil and 
+      self.count_include_pad or true
+   if self.count_include_pad then
+      self.mode = 'CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING'
+   else 
+      error'This mode is untested in cudnn'
+      self.mode = 'CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING'
+   end
+   return parent.updateOutput(self, input)
 end
 
 function SpatialAveragePooling:__tostring__()