Add SpatialRadialMatching

Author: Michael 'myrhev' Mathieu

SpatialPadding.lua

@@ -24,6 +24,9 @@ function SpatialPadding:__init(pad_l, pad_r, pad_t, pad_b, y_dim, x_dim)
 end
 
 function SpatialPadding:updateOutput(input)
+   if self.output:type() ~= input:type() then
+      self.output = input.new()
+   end
    self.x_dim = self.x_dim or 3
    self.y_dim = self.y_dim or 2
    local h = input:size(self.y_dim) + self.pad_t + self.pad_b

SpatialRadialMatching.lua

@@ -0,0 +1,33 @@
+local SpatialRadialMatching, parent = torch.class('nn.SpatialRadialMatching', 'nn.Module')
+
+function SpatialRadialMatching:__init(maxh)
+   -- If full_output is false, output is computed on elements of the first input
+   -- for which all the possible corresponding elements exist in the second input
+   -- In addition, if full_output is set to false, the pixel (1,1) of the first input
+   -- is supposed to correspond to the pixel (maxh/2, maxw/2) of the second one
+   parent.__init(self)
+   self.maxh = maxh
+   self.gradInput1 = torch.Tensor()
+   self.gradInput2 = torch.Tensor()
+end
+
+function SpatialRadialMatching:updateOutput(input)
+   -- input is a table of 2 inputs, each one being KxHxW
+   -- if not full_output, the 1st one is KxH1xW1 where H1 <= H-maxh+1, W1 <= W-maxw+1
+   self.output:resize(input[1]:size(2), input[1]:size(3), self.maxh)
+   --if input[3] == nil then
+   --   input[3] = torch.LongTensor(input[1]:size(2), input[1]:size(3)):fill(1)
+   --end
+   --input[1].nn.SpatialRadialMatching_updateOutput(self, input[1], input[2], input[3])
+   input[1].nn.SpatialRadialMatching_updateOutput(self, input[1], input[2])
+   return self.output
+end
+
+function SpatialRadialMatching:updateGradInput(input, gradOutput)
+   self.gradInput1:resize(input[1]:size()):zero()
+   self.gradInput2:resize(input[2]:size()):zero()
+   --input[1].nn.SpatialRadialMatching_updateGradInput(self,input[1],input[2],gradOutput,input[3])
+input[1].nn.SpatialRadialMatching_updateGradInput(self,input[1],input[2],gradOutput)
+   self.gradInput = {self.gradInput1, self.gradInput2}
+   return self.gradInput
+end

generic/SpatialMatching.c

@@ -89,7 +89,7 @@ static int nn_(SpatialMatching_updateOutput)(lua_State *L)
     }
     */
   } else {
-    //#pragma omp parallel for private(y1,x1,x2,y2,k,dist)
+#pragma omp parallel for private(y1,x1,x2,y2,k,dist)
     for (y1 = 0; y1 < iheight; y1++) {
       for (x1 = 0; x1 < iwidth; x1++) {
 	for (y2 = y1; y2 < y1+maxh; y2++) {

generic/SpatialRadialMatching.c

@@ -0,0 +1,125 @@
+#ifndef TH_GENERIC_FILE
+#define TH_GENERIC_FILE "generic/SpatialRadialMatching.c"
+#else
+
+#define square(x) ((x)*(x))
+#define max(x,y) (((x)>(y)) ? (x) : (y))
+#define min(x,y) (((x)>(y)) ? (y) : (x))
+
+static int nn_(SpatialRadialMatching_updateOutput)(lua_State *L)
+{
+  // get all params
+  THTensor *input1  = luaT_checkudata(L, 2, torch_(Tensor_id));
+  THTensor *input2  = luaT_checkudata(L, 3, torch_(Tensor_id));
+  //THLongTensor *mask= luaT_checkudata(L, 4, luaT_checktypename2id(L, "torch.LongTensor"));
+  int maxh          = luaT_getfieldcheckint(L, 1, "maxh");
+  THTensor *output  = luaT_getfieldcheckudata(L, 1, "output", torch_(Tensor_id));
+
+  // dims
+  int iwidth = input1->size[2];
+  int iheight = input1->size[1];
+  int ichannels = input1->size[0];
+
+  // get strides
+  long *i1s = input1->stride;
+  long *i2s = input2->stride;
+  //long *ms  = mask  ->stride;
+  long *os  = output->stride;
+
+  // get pointers
+  real *input1_p = THTensor_(data)(input1);
+  real *input2_p = THTensor_(data)(input2);
+  //long *mask_p   = THLongTensor_data(mask);
+  real *output_p = THTensor_(data)(output);
+
+  // compute output
+  int x1,y1,y2,k;
+  real dist;
+#pragma omp parallel for private(y1,x1,y2,k,dist)
+  for (y1 = 0; y1 < iheight; y1++) {
+    for (x1 = 0; x1 < iwidth; x1++) {
+      //if (mask_p[y1*ms[0] + x1*ms[1]]) {
+	for (y2 = y1; y2 < y1+maxh; y2++) {
+	  dist = 0.0f;
+	  for (k = 0; k < ichannels; k++)
+	    dist += square(  input1_p[k*i1s[0] + y1*i1s[1] + x1*i1s[2]]
+			     - input2_p[k*i2s[0] + y2*i2s[1] + x1*i2s[2]]);
+	  output_p[(y2-y1)*os[2] + y1*os[0] + x1*os[1]] = dist;
+	}
+	//}
+    }
+  }
+  
+  // done
+  return 0;
+}
+
+static int nn_(SpatialRadialMatching_updateGradInput)(lua_State *L)
+{
+  // get all params
+  THTensor*     input1 = luaT_checkudata(L, 2, torch_(Tensor_id));
+  THTensor*     input2 = luaT_checkudata(L, 3, torch_(Tensor_id));
+  THTensor* gradOutput = luaT_checkudata(L, 4, torch_(Tensor_id));
+  //THLongTensor*   mask = luaT_checkudata(L, 5, luaT_checktypename2id(L, "torch.LongTensor"));
+  THTensor* gradInput1 = luaT_getfieldcheckudata(L, 1, "gradInput1", torch_(Tensor_id));
+  THTensor* gradInput2 = luaT_getfieldcheckudata(L, 1, "gradInput2", torch_(Tensor_id));
+  int             maxh = luaT_getfieldcheckint(L, 1, "maxh");
+
+  // dims
+  int iwidth    = input1->size[2];
+  int iheight   = input1->size[1];
+  int ichannels = input1->size[0];
+
+  // get strides
+  long* i1s  = input1->stride;
+  long* i2s  = input2->stride;
+  long* gi1s = gradInput1->stride;
+  long* gi2s = gradInput2->stride;
+  long* gos  = gradOutput->stride;
+  //long* ms   = mask->stride;
+  
+  // get pointers
+  real* input1_p     = THTensor_(data)(input1);
+  real* input2_p     = THTensor_(data)(input2);
+  real* gradInput1_p = THTensor_(data)(gradInput1);
+  real* gradInput2_p = THTensor_(data)(gradInput2);
+  real* gradOutput_p = THTensor_(data)(gradOutput);
+  //long* mask_p       = THLongTensor_data(mask);
+  
+  // compute gradients
+  int x1, y1, y2, k;
+  real partial_d;
+  for (y1 = 0; y1 < iheight; y1++) {
+    for (x1 = 0; x1 < iwidth; x1++) {
+      // if (mask_p[y1*ms[0] + x1*ms[1]]) {
+	for (y2 = y1; y2 < y1+maxh; y2++) {
+	  for (k = 0; k < ichannels; k++) {
+	    partial_d = 2.0f*(  input1_p[k*i1s[0] + y1*i1s[1] + x1*i1s[2]]
+				- input2_p[k*i2s[0] + y2*i2s[1] + x1*i2s[2]]);
+	    partial_d *= gradOutput_p[(y2-y1)*gos[2]+y1*gos[0]+x1*gos[1]];
+	    gradInput1_p[k*gi1s[0] + y1*gi1s[1] + x1*gi1s[2]] += partial_d;
+	    gradInput2_p[k*gi2s[0] + y2*gi2s[1] + x1*gi2s[2]] -= partial_d;
+	  }
+	}
+	//}
+    }
+  }
+
+  // done
+  return 0;
+}
+
+static const struct luaL_Reg nn_(SpatialRadialMatching__) [] = {
+  {"SpatialRadialMatching_updateOutput", nn_(SpatialRadialMatching_updateOutput)},
+  {"SpatialRadialMatching_updateGradInput", nn_(SpatialRadialMatching_updateGradInput)},
+  {NULL, NULL}
+};
+
+static void nn_(SpatialRadialMatching_init)(lua_State *L)
+{
+  luaT_pushmetaclass(L, torch_(Tensor_id));
+  luaT_registeratname(L, nn_(SpatialRadialMatching__), "nn");
+  lua_pop(L,1);
+}
+
+#endif

init.c

@@ -57,6 +57,9 @@ static const void* torch_DoubleTensor_id = NULL;
 #include "generic/SpatialMatching.c"
 #include "THGenerateFloatTypes.h"
 
+#include "generic/SpatialRadialMatching.c"
+#include "THGenerateFloatTypes.h"
+
 #include "generic/DataSetLabelMe.c"
 #include "THGenerateFloatTypes.h"
 
@@ -81,6 +84,7 @@ DLL_EXPORT int luaopen_libnnx(lua_State *L)
   nn_FloatSpatialClassNLLCriterion_init(L);
   nn_FloatSpatialGraph_init(L);
   nn_FloatSpatialMatching_init(L);
+  nn_FloatSpatialRadialMatching_init(L);
   nn_FloatDataSetLabelMe_init(L);
 
   nn_DoubleSpatialLinear_init(L);
@@ -99,6 +103,7 @@ DLL_EXPORT int luaopen_libnnx(lua_State *L)
   nn_DoubleSpatialClassNLLCriterion_init(L);
   nn_DoubleSpatialGraph_init(L);
   nn_DoubleSpatialMatching_init(L);
+  nn_DoubleSpatialRadialMatching_init(L);
   nn_DoubleDataSetLabelMe_init(L);
 
   return 1;

init.lua

@@ -73,6 +73,7 @@ torch.include('nnx', 'SpatialFovea.lua')
 torch.include('nnx', 'SpatialPyramid.lua')
 torch.include('nnx', 'SpatialGraph.lua')
 torch.include('nnx', 'SpatialMatching.lua')
+torch.include('nnx', 'SpatialRadialMatching.lua')
 torch.include('nnx', 'SpatialMaxSampling.lua')
 torch.include('nnx', 'SpatialColorTransform.lua')
 torch.include('nnx', 'SpatialConvolutionSparse.lua')