Add Spatial Dilated Max Pooling

new file:   SpatialDilatedMaxPooling.lua
modified:   SpatialMaxPooling.lua
modified:   init.lua
modified:   lib/THNN/generic/SpatialMaxPooling.c
modified:   test.lua
Updated THNN
modified:   THNN.h

Removed unnecessary function calls
modified:   SpatialDilatedMaxPooling.lua

6 files changed, 144 insertions, 18 deletions

diff --git a/SpatialDilatedMaxPooling.lua b/SpatialDilatedMaxPooling.lua
new file mode 100644
index 0000000..929459c
--- /dev/null
+++ b/SpatialDilatedMaxPooling.lua
@@ -0,0 +1,62 @@
+local THNN = require 'nn.THNN'
+local SpatialDilatedMaxPooling, parent = torch.class('nn.SpatialDilatedMaxPooling', 'nn.SpatialMaxPooling')
+
+function SpatialDilatedMaxPooling:__init(kW, kH, dW, dH, padW, padH, dilationW, dilationH)
+   parent.__init(self, kW, kH, dW, dH, padW, padH)
+
+   self.dilationW = dilationW or 1
+   self.dilationH = dilationH or 1
+end
+
+function SpatialDilatedMaxPooling:updateOutput(input)
+   self.indices = self.indices or input.new()
+
+   local dims = input:dim()
+   self.iheight = input:size(dims-1)
+   self.iwidth = input:size(dims)
+
+   input.THNN.SpatialMaxPooling_updateOutput(
+      input:cdata(),
+      self.output:cdata(),
+      self.indices:cdata(),
+      self.kW, self.kH,
+      self.dW, self.dH,
+      self.padW, self.padH,
+      self.dilationW, self.dilationH,
+      self.ceil_mode
+   )
+   return self.output
+end
+
+function SpatialDilatedMaxPooling:updateGradInput(input, gradOutput)
+   input.THNN.SpatialMaxPooling_updateGradInput(
+      input:cdata(),
+      gradOutput:cdata(),
+      self.gradInput:cdata(),
+      self.indices:cdata(),
+      self.kW, self.kH,
+      self.dW, self.dH,
+      self.padW, self.padH,
+      self.dilationW, self.dilationH,
+      self.ceil_mode
+   )
+   return self.gradInput
+end
+
+function SpatialDilatedMaxPooling:__tostring__()
+   local s =  string.format('%s(%dx%d, %d,%d', torch.type(self),
+                            self.kW, self.kH, self.dW, self.dH)
+   if (self.padW or self.padH) and (self.padW ~= 0 or self.padH ~= 0) then
+      s = s .. ', ' .. self.padW .. ','.. self.padH
+   end
+   s = s .. ', ' .. self.dilationW .. ',' .. self.dilationH
+   s = s .. ')'
+   return s
+end
+
+function SpatialDilatedMaxPooling:clearState()
+   if self.indices then
+      self.indices:set()
+   end
+   return parent.clearState(self)
+end
diff --git a/SpatialMaxPooling.lua b/SpatialMaxPooling.lua
index 8475b13..c05a876 100644
--- a/SpatialMaxPooling.lua
+++ b/SpatialMaxPooling.lua
@@ -46,6 +46,7 @@ function SpatialMaxPooling:updateOutput(input)
       self.kW, self.kH,
       self.dW, self.dH,
       self.padW, self.padH,
+      1, 1,
       self.ceil_mode
    )
    return self.output
@@ -60,6 +61,7 @@ function SpatialMaxPooling:updateGradInput(input, gradOutput)
       self.kW, self.kH,
       self.dW, self.dH,
       self.padW, self.padH,
+      1, 1,
       self.ceil_mode
    )
    return self.gradInput
diff --git a/init.lua b/init.lua
index fe36d8c..a9c68da 100644
--- a/init.lua
+++ b/init.lua
@@ -101,6 +101,7 @@ require('nn.SpatialConvolutionMap')
 require('nn.SpatialDilatedConvolution')
 require('nn.SpatialSubSampling')
 require('nn.SpatialMaxPooling')
+require('nn.SpatialDilatedMaxPooling')
 require('nn.SpatialMaxUnpooling')
 require('nn.SpatialFractionalMaxPooling')
 require('nn.SpatialLPPooling')
diff --git a/lib/THNN/generic/SpatialMaxPooling.c b/lib/THNN/generic/SpatialMaxPooling.c
index 829f3f0..3daef1d 100644
--- a/lib/THNN/generic/SpatialMaxPooling.c
+++ b/lib/THNN/generic/SpatialMaxPooling.c
@@ -16,7 +16,10 @@ static void THNN_(SpatialMaxPooling_updateOutput_frame)(
           int dW,
           int dH,
           int padW,
-          int padH)
+          int padH,
+          int dilationW,
+          int dilationH
+          )
 {
   long k;
 #pragma omp parallel for private(k)
@@ -31,10 +34,12 @@ static void THNN_(SpatialMaxPooling_updateOutput_frame)(
       {
         long hstart = i * dH - padH;
         long wstart = j * dW - padW;
-        long hend = fminf(hstart + kH, iheight);
-        long wend = fminf(wstart + kW, iwidth);
-        hstart = fmaxf(hstart, 0);
-        wstart = fmaxf(wstart, 0);
+        long hend = fminf(hstart + (kH - 1) * dilationH + 1, iheight);
+        long wend = fminf(wstart + (kW - 1) * dilationW + 1, iwidth);
+        while(hstart < 0)
+          hstart += dilationH;
+        while(wstart < 0)
+          wstart += dilationW;
 
         /* local pointers */
         real *op = output_p  + k*owidth*oheight + i*owidth + j;
@@ -45,9 +50,9 @@ static void THNN_(SpatialMaxPooling_updateOutput_frame)(
         real maxval = -THInf;
         long tcntr = 0;
         long x,y;
-        for(y = hstart; y < hend; y++)
+        for(y = hstart; y < hend; y += dilationH)
         {
-          for(x = wstart; x < wend; x++)
+          for(x = wstart; x < wend; x += dilationW)
           {
             tcntr = y*iwidth + x;
             real val = *(ip + tcntr);
@@ -80,6 +85,8 @@ void THNN_(SpatialMaxPooling_updateOutput)(
           int dH,
           int padW,
           int padH,
+          int dilationW,
+          int dilationH,
           bool ceil_mode)
 {
   int dimw = 2;
@@ -104,24 +111,27 @@ void THNN_(SpatialMaxPooling_updateOutput)(
     dimh++;
   }
   THArgCheck(input->size[dimw] >= kW - padW && input->size[dimh] >= kH - padH, 2, "input image smaller than kernel size");
-
   THArgCheck(kW/2 >= padW && kH/2 >= padH, 2, "pad should be smaller than half of kernel size");
-
+  
   /* sizes */
   nslices = input->size[dimh-1];
   iheight = input->size[dimh];
   iwidth = input->size[dimw];
   if (ceil_mode)
   {
-    oheight = (long)(ceil((float)(iheight - kH + 2*padH) / dH)) + 1;
-    owidth  = (long)(ceil((float)(iwidth  - kW + 2*padW) / dW)) + 1;
+    oheight = (long)(ceil((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
+    owidth  = (long)(ceil((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
   }
   else
   {
-    oheight = (long)(floor((float)(iheight - kH + 2*padH) / dH)) + 1;
-    owidth  = (long)(floor((float)(iwidth  - kW + 2*padW) / dW)) + 1;
+    oheight = (long)(floor((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
+    owidth  = (long)(floor((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
   }
 
+  if (owidth < 1 || oheight < 1)
+    THError("Given input size: (%dx%dx%d). Calculated output size: (%dx%dx%d). Output size is too small",
+            nslices,iheight,iwidth,nslices,oheight,owidth);
+
   if (padW || padH)
   {
     // ensure that the last pooling starts inside the image
@@ -151,7 +161,9 @@ void THNN_(SpatialMaxPooling_updateOutput)(
                                               iwidth, iheight,
                                               owidth, oheight,
                                               kW, kH, dW, dH,
-                                              padW, padH);
+                                              padW, padH,
+                                              dilationW, dilationH
+                                              );
   }
   else
   {
@@ -174,7 +186,9 @@ void THNN_(SpatialMaxPooling_updateOutput)(
                                                 iwidth, iheight,
                                                 owidth, oheight,
                                                 kW, kH, dW, dH,
-                                                padW, padH);
+                                                padW, padH,
+                                                dilationW, dilationH
+                                                );
     }
   }
 
@@ -229,6 +243,8 @@ void THNN_(SpatialMaxPooling_updateGradInput)(
           int dH,
           int padW,
           int padH,
+          int dilationW,
+          int dilationH,
           bool ceil_mode)
 {
   int dimw = 2;
diff --git a/lib/THNN/generic/THNN.h b/lib/THNN/generic/THNN.h
index 974f56c..7ad6f70 100644
--- a/lib/THNN/generic/THNN.h
+++ b/lib/THNN/generic/THNN.h
@@ -840,6 +840,7 @@ TH_API void THNN_(SpatialMaxPooling_updateOutput)(
           int kW, int kH,
           int dW, int dH,
           int padW, int padH,
+          int dilationW, int dilationH,
           bool ceil_mode);
 TH_API void THNN_(SpatialMaxPooling_updateGradInput)(
           THNNState *state,
@@ -850,6 +851,7 @@ TH_API void THNN_(SpatialMaxPooling_updateGradInput)(
           int kW, int kH,
           int dW, int dH,
           int padW, int padH,
+          int dilationW, int dilationH,
           bool ceil_mode);
 
 TH_API void THNN_(SpatialMaxUnpooling_updateOutput)(
diff --git a/test.lua b/test.lua
index 8327538..e288e25 100644
--- a/test.lua
+++ b/test.lua
@@ -1676,9 +1676,9 @@ function nntest.LogSoftmax()
    local gradInput2 = layer:backward(input, gradOutput):clone()
 
    mytester:assertlt(gradInput1:add(-1, gradInput2):abs():max(),
-		     1e-10,
-		     torch.typename(layer)
-			.. ' non-contiguous gradOutput check')
+           1e-10,
+           torch.typename(layer)
+         .. ' non-contiguous gradOutput check')
 
 
 
@@ -3210,6 +3210,49 @@ function nntest.SpatialMaxUnpooling()
   end
 end
 
+function nntest.SpatialDilatedMaxPooling()
+   for _,ceil_mode in pairs({true,false}) do
+      local from = math.random(1,5)
+      local ki = math.random(1,4)
+      local kj = math.random(1,4)
+      local si = math.random(1,3)
+      local sj = math.random(1,3)
+      local outi = math.random(4,5)
+      local outj = math.random(4,5)
+      local padW = math.min(math.random(0,1),math.floor(ki/2))
+      local padH =  math.min(math.random(0,1),math.floor(kj/2))
+      local dilationW = math.random(1,5)
+      local dilationH = math.random(1,5)
+      local ini = (outi-1)*si+(dilationW*(ki-1)+1)-2*padW
+      local inj = (outj-1)*sj+(dilationH*(kj-1)+1)-2*padH
+
+      local ceil_string = ceil_mode and 'ceil' or 'floor'
+      local module = nn.SpatialDilatedMaxPooling(ki,kj,si,sj,padW,padH,dilationW, dilationH)
+      if ceil_mode then module:ceil() else module:floor() end
+      local input = torch.rand(from,inj,ini)
+
+      local err = jac.testJacobian(module, input)
+      mytester:assertlt(err, precision, 'error '..ceil_string..' mode on state ')
+
+      local ferr, berr = jac.testIO(module, input)
+      mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err ')
+      mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err ')
+
+      -- batch
+      local nbatch = math.random(2,5)
+      input = torch.rand(nbatch,from,inj,ini)
+      module = nn.SpatialDilatedMaxPooling(ki,kj,si,sj,padW,padH,dilationW,dilationH)
+      if ceil_mode then module:ceil() else module:floor() end
+
+      local err = jac.testJacobian(module, input)
+      mytester:assertlt(err, precision, 'error '..ceil_string..' mode on state (Batch)')
+
+      local ferr, berr = jac.testIO(module, input)
+      mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err (Batch) ')
+      mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err (Batch) ')
+  end
+end
+
 function nntest.SpatialFractionalMaxPooling()
     local batch = math.random(1, 3)
     local plane = math.random(1, 3)