VolumetricFractionalMaxPooling like Spatial...

3 files changed, 299 insertions, 0 deletions

diff --git a/lib/THNN/generic/THNN.h b/lib/THNN/generic/THNN.h
index 4420962..076c178 100644
--- a/lib/THNN/generic/THNN.h
+++ b/lib/THNN/generic/THNN.h
@@ -1131,6 +1131,23 @@ TH_API void THNN_(VolumetricConvolutionMM_accGradParameters)(
           int pT, int pW, int pH,
           real scale);
 
+TH_API void THNN_(VolumetricFractionalMaxPooling_updateOutput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *output,
+          int outputT, int outputW, int outputH,
+          int poolSizeT, int poolSizeW, int poolSizeH,
+          THIndexTensor *indices,
+          THTensor *randomSamples);
+TH_API void THNN_(VolumetricFractionalMaxPooling_updateGradInput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *gradOutput,
+          THTensor *gradInput,
+          int outputT, int outputW, int outputH,
+          int poolSizeT, int poolSizeW, int poolSizeH,
+          THIndexTensor *indices);
+
 TH_API void THNN_(VolumetricFullConvolution_updateOutput)(
           THNNState *state,         // library state
           THTensor *input,          // 4D or 5D (batch) tensor
diff --git a/lib/THNN/generic/VolumetricFractionalMaxPooling.c b/lib/THNN/generic/VolumetricFractionalMaxPooling.c
new file mode 100644
index 0000000..e7b7208
--- /dev/null
+++ b/lib/THNN/generic/VolumetricFractionalMaxPooling.c
@@ -0,0 +1,279 @@
+#ifndef TH_GENERIC_FILE
+#define TH_GENERIC_FILE "generic/VolumetricFractionalMaxPooling.c"
+#else
+
+static long* THNN_(VolumetricFractionalMaxPooling_generateIntervals)(
+  real sample,
+  long inputSize,
+  long outputSize,
+  int poolSize) {
+  real alpha = (real) (inputSize - poolSize) / (real) (outputSize - 1);
+  long* sequence = (long*) THAlloc(sizeof(long) * outputSize);
+
+  long i;
+  for (i = 0; i < outputSize - 1; ++i) {
+    sequence[i] =
+      (long) ((i + sample) * alpha) - (long) (sample * alpha);
+  }
+  sequence[outputSize - 1] = inputSize - poolSize;
+
+  return sequence;
+}
+
+static void THNN_(VolumetricFractionalMaxPooling_updateOutput_frame)(
+  real* input,
+  real* output,
+  THIndex_t* indices,
+  real* randomSamples,
+  long numPlanes,
+  long inputT, long inputW, long inputH,
+  long outputT, long outputW, long outputH,
+  int poolSizeT, int poolSizeW, int poolSizeH) {
+  long plane;
+#pragma omp parallel for private(plane)
+  for (plane = 0; plane < numPlanes; ++plane) {
+    /* each plane contains 3 random samples, one for T, one for W, and one for H */
+    real* randomSamplesForPlane = randomSamples + plane * 3;
+
+    /* Generate interval sequence */
+    long* sequenceT =
+      THNN_(VolumetricFractionalMaxPooling_generateIntervals)(
+        randomSamplesForPlane[0], inputT, outputT, poolSizeT);
+    long* sequenceW =
+      THNN_(VolumetricFractionalMaxPooling_generateIntervals)(
+        randomSamplesForPlane[1], inputW, outputW, poolSizeW);
+    long* sequenceH =
+      THNN_(VolumetricFractionalMaxPooling_generateIntervals)(
+        randomSamplesForPlane[2], inputH, outputH, poolSizeH);
+
+    /* loop over output */
+    long h, w, t;
+
+    real* inputForPlane = input + plane * inputT * inputW * inputH;
+    real* outputForPlane = output + plane * outputT * outputW * outputH;
+    THIndex_t* indicesForPlane = indices + plane * outputT * outputW * outputH;
+
+    for (h = 0; h < outputH; ++h) {
+      long inputHStart = sequenceH[h];
+
+      for (w = 0; w < outputW; ++w) {
+        long inputWStart = sequenceW[w];
+
+        for (t = 0; t < outputT; ++t) {
+          long inputTStart = sequenceT[t];
+
+          real maxVal = -THInf;
+          long maxIndex = -1;
+
+          long h2, w2, t2;
+          for (h2 = inputHStart; h2 < inputHStart + poolSizeH; ++h2) {
+            for (w2 = inputWStart; w2 < inputWStart + poolSizeW; ++w2) {
+              for (t2 = inputTStart; t2 < inputTStart + poolSizeT; ++t2) {
+                THAssert(h2 >= 0 && h2 < inputH);
+                THAssert(w2 >= 0 && w2 < inputW);
+                THAssert(t2 >= 0 && t2 < inputT);
+
+                long planeIndex = h2 * inputW * inputT + w2 * inputT + t2;
+                real val = inputForPlane[planeIndex];
+                if (val > maxVal) {
+                  maxVal = val;
+                  maxIndex = planeIndex;
+                }
+              }
+            }
+          }
+
+          THAssert(maxVal != -THInf);
+          THAssert(maxIndex != -1);
+
+          outputForPlane[h * outputW * outputT + w * outputT + t] = maxVal;
+          /* +1 to lua index */
+          indicesForPlane[h * outputW * outputT + w * outputT + t] = maxIndex + TH_INDEX_BASE;
+        }
+      }
+    }
+
+    THFree(sequenceT);
+    THFree(sequenceW);
+    THFree(sequenceH);
+  }
+}
+
+void THNN_(VolumetricFractionalMaxPooling_updateOutput)(
+    THNNState *state,
+    THTensor *input,
+    THTensor *output,
+    int outputT, int outputW, int outputH,
+    int poolSizeT, int poolSizeW, int poolSizeH,
+    THIndexTensor *indices,
+    THTensor *randomSamples) {
+
+  long numBatch = 1;
+  int planeDim = 0;
+  int heightDim = 1;
+  int widthDim = 2;
+  int timeDim = 3;
+
+  long numInputDims = THTensor_(nDimension)(input);
+  THNN_ARGCHECK(numInputDims == 4 || numInputDims == 5, 2, input,
+		"4D or 5D (batch mode) tensor expected for input, but got: %s");
+
+  if (numInputDims == 5) {
+    numBatch = THTensor_(size)(input, 0);
+    planeDim++;
+    heightDim++;
+    widthDim++;
+    timeDim++;
+  }
+
+  /* sizes */
+  long numPlanes = THTensor_(size)(input, planeDim);
+  long inputH = THTensor_(size)(input, heightDim);
+  long inputW = THTensor_(size)(input, widthDim);
+  long inputT = THTensor_(size)(input, timeDim);
+
+  THArgCheck(outputH + poolSizeH - 1 < inputH, 9,
+             "poolSizeH (%d) too large relative to input height (%d)",
+	     poolSizeH, inputH);
+  THArgCheck(outputW + poolSizeW - 1 < inputW, 8,
+             "poolSizeW (%d) too large relative to input width (%d)",
+	     poolSizeW, inputW);
+  THArgCheck(outputT + poolSizeT - 1 < inputT, 7,
+             "poolSizeW (%d) too large relative to input width (%d)",
+	     poolSizeT, inputT);
+
+  /* get contiguous input */
+  input = THTensor_(newContiguous)(input);
+
+  if (numInputDims == 4) {
+    /* resize output */
+    THTensor_(resize4d)(output, numPlanes, outputH, outputW, outputT);
+    /* indices will contain the locations for each output point */
+    THIndexTensor_(resize4d)(indices, numPlanes, outputH, outputW, outputT);
+
+    THNN_(VolumetricFractionalMaxPooling_updateOutput_frame)(
+      THTensor_(data)(input),
+      THTensor_(data)(output),
+      THIndexTensor_(data)(indices),
+      THTensor_(data)(randomSamples),
+      numPlanes, inputT, inputW, inputH,
+      outputT, outputW, outputH, poolSizeT, poolSizeW, poolSizeH);
+  } else {
+    THTensor_(resize5d)(output, numBatch, numPlanes, outputH, outputW, outputT);
+    /* indices will contain the locations for each output point */
+    THIndexTensor_(resize5d)(indices, numBatch, numPlanes, outputH, outputW, outputT);
+
+    long batch;
+#pragma omp parallel for private(batch)
+    for (batch = 0; batch < numBatch; ++batch) {
+      THNN_(VolumetricFractionalMaxPooling_updateOutput_frame)(
+        THTensor_(data)(input) + batch * numPlanes * inputH * inputW * inputT,
+        THTensor_(data)(output) + batch * numPlanes * outputH * outputW * outputT,
+        THIndexTensor_(data)(indices) + batch * numPlanes * outputH * outputW * outputT,
+        THTensor_(data)(randomSamples) + batch * numPlanes * 3,
+        numPlanes, inputT, inputW, inputH,
+        outputT, outputW, outputH, poolSizeT, poolSizeW, poolSizeH);
+    }
+  }
+
+  /* cleanup */
+  THTensor_(free)(input);
+}
+
+static void THNN_(VolumetricFractionalMaxPooling_updateGradInput_frame)(
+  real* gradInput,
+  real* gradOutput,
+  THIndex_t* indices,
+  long numPlanes,
+  long inputT, long inputW, long inputH,
+  long outputT, long outputW, long outputH) {
+  long plane;
+#pragma omp parallel for private(plane)
+  for (plane = 0; plane < numPlanes; plane++) {
+    real* gradInputForPlane = gradInput + plane * inputT * inputW * inputH;
+    real* gradOutputForPlane = gradOutput + plane * outputT * outputW * outputH;
+    THIndex_t* indicesForPlane = indices + plane * outputT * outputW * outputH;
+
+    long h, w, t;
+    for (h = 0; h < outputH; ++h) {
+      for (w = 0; w < outputW; ++w) {
+        for (t = 0; t < outputT; ++t) {
+          long outputIndex = h * outputW * outputT + w * outputT+ t;
+          long index = indicesForPlane[outputIndex] - TH_INDEX_BASE;
+          THAssert(index >= 0 && index < inputT * inputW * inputH);
+
+          gradInputForPlane[index] += gradOutputForPlane[outputIndex];
+        }
+      }
+    }
+  }
+}
+
+void THNN_(VolumetricFractionalMaxPooling_updateGradInput)(
+    THNNState *state,
+    THTensor *input,
+    THTensor *gradOutput,
+    THTensor *gradInput,
+    int outputT, int outputW, int outputH,
+    int poolSizeT, int poolSizeW, int poolSizeH,
+    THIndexTensor *indices) {
+
+  long numBatch = 1;
+  int planeDim = 0;
+  int heightDim = 1;
+  int widthDim = 2;
+  int timeDim = 3;
+
+  long numInputDims = THTensor_(nDimension)(input);
+  if (numInputDims == 5) {
+    numBatch = THTensor_(size)(input, 0);
+    planeDim = 1;
+    heightDim++;
+    widthDim++;
+    timeDim++;
+  }
+
+  /* sizes */
+  long numPlanes = THTensor_(size)(input, planeDim);
+  long inputH = THTensor_(size)(input, heightDim);
+  long inputW = THTensor_(size)(input, widthDim);
+  long inputT = THTensor_(size)(input, timeDim);
+
+  THArgCheck(outputT == THTensor_(size)(gradOutput, timeDim), 3,
+             "gradOutput time unexpected");
+  THArgCheck(outputW == THTensor_(size)(gradOutput, widthDim), 3,
+             "gradOutput width unexpected");
+  THArgCheck(outputH == THTensor_(size)(gradOutput, heightDim), 3,
+             "gradOutput height unexpected");
+
+  /* get contiguous gradOutput */
+  gradOutput = THTensor_(newContiguous)(gradOutput);
+
+  /* resize */
+  THTensor_(resizeAs)(gradInput, input);
+  THTensor_(zero)(gradInput);
+
+  /* backprop */
+  if (numInputDims == 4) {
+    THNN_(VolumetricFractionalMaxPooling_updateGradInput_frame)(
+      THTensor_(data)(gradInput),
+      THTensor_(data)(gradOutput),
+      THIndexTensor_(data)(indices),
+      numPlanes, inputT, inputW, inputH, outputT, outputW, outputH);
+  } else {
+    long batch;
+#pragma omp parallel for private(batch)
+    for (batch = 0; batch < numBatch; ++batch) {
+      THNN_(VolumetricFractionalMaxPooling_updateGradInput_frame)(
+        THTensor_(data)(gradInput) + batch * numPlanes * inputH * inputW * inputT,
+        THTensor_(data)(gradOutput) + batch * numPlanes * outputH * outputW * outputT,
+        THIndexTensor_(data)(indices) + batch * numPlanes * outputH * outputW * outputT,
+        numPlanes, inputT, inputW, inputH, outputT, outputW, outputH);
+    }
+  }
+
+  /* cleanup */
+  THTensor_(free)(gradOutput);
+}
+
+#endif
diff --git a/lib/THNN/init.c b/lib/THNN/init.c
index 990775d..8db59bd 100644
--- a/lib/THNN/init.c
+++ b/lib/THNN/init.c
@@ -245,6 +245,9 @@
 #include "generic/VolumetricDilatedMaxPooling.c"
 #include "THGenerateFloatTypes.h"
 
+#include "generic/VolumetricFractionalMaxPooling.c"
+#include "THGenerateFloatTypes.h"
+
 #include "generic/VolumetricMaxUnpooling.c"
 #include "THGenerateFloatTypes.h"