Merge pull request #123 from torch/volconvbatch

VolumetricConvolution batch mode + test

2 files changed, 150 insertions, 37 deletions

diff --git a/generic/VolumetricConvolution.c b/generic/VolumetricConvolution.c
index feeaf05..a7960dc 100644
--- a/generic/VolumetricConvolution.c
+++ b/generic/VolumetricConvolution.c
@@ -13,22 +13,31 @@ static int nn_(VolumetricConvolution_updateOutput)(lua_State *L)
   THTensor *bias = luaT_getfieldcheckudata(L, 1, "bias", torch_Tensor);
   THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_Tensor);
 
-  luaL_argcheck(L, input->nDimension == 4, 2, "4D tensor expected");
-
-  {
-    long nOutputPlane = weight->size[0];
-    long kT           = weight->size[2];
-    long kH           = weight->size[3];
-    long kW           = weight->size[4];
-    long inputDepth   = input->size[1];
-    long inputHeight  = input->size[2];
-    long inputWidth   = input->size[3];
-    long outputDepth  = (inputDepth - kT) / dT + 1;
-    long outputWidth  = (inputWidth - kW) / dW + 1;
-    long outputHeight = (inputHeight - kH) / dH + 1;
-    THTensor *outn = THTensor_(new)();
-    long i;
+  luaL_argcheck(L, input->nDimension == 4 || input->nDimension == 5, 
+		2, "4D or 5D (batch-mode) tensor expected");
+  int dimt = 1;
+  int dimh = 2;
+  int dimw = 3;
+
+  if (input->nDimension == 5) {
+    dimt++;
+    dimh++;
+    dimw++;
+  }
 
+  long nOutputPlane = weight->size[0];
+  long kT           = weight->size[2];
+  long kH           = weight->size[3];
+  long kW           = weight->size[4];
+  long inputDepth   = input->size[dimt];
+  long inputHeight  = input->size[dimh];
+  long inputWidth   = input->size[dimw];
+  long outputDepth  = (inputDepth - kT) / dT + 1;
+  long outputWidth  = (inputWidth - kW) / dW + 1;
+  long outputHeight = (inputHeight - kH) / dH + 1;
+  THTensor *outn = THTensor_(new)();
+  long i,j;
+  if (input->nDimension == 4) { /* non-batch mode */
     THTensor_(resize4d)(output, nOutputPlane, outputDepth, outputHeight, outputWidth);
   
     /* add bias */
@@ -37,18 +46,42 @@ static int nn_(VolumetricConvolution_updateOutput)(lua_State *L)
       THTensor_(fill)(outn, THTensor_(get1d)(bias, i));
     }
 
-    THTensor_(free)(outn);
-
     /* do convolutions */
     THTensor_(conv3Dmv)(output, 1.0, 1.0, input, weight, dT, dH, dW, "V", "X");
-  }
+  } else { /* batch mode */
+    long nBatch = input->size[0];
+    THTensor_(resize5d)(output, nBatch, nOutputPlane, 
+			outputDepth, outputHeight, outputWidth);
+    THTensor *inb = THTensor_(new)();
+    THTensor *outb = THTensor_(new)();
+
+#pragma omp parallel for private(j)
+    for (j=0; j<nBatch; j++) { /* loop over batches */
+      THTensor_(select)(inb,input,0,j);
+      THTensor_(select)(outb,output,0,j);
+	
+      /* add bias */
+      for (i=0; i<bias->size[0]; i++) {
+	THTensor_(select)(outn,outb,0,i);
+	THTensor_(fill)(outn, THTensor_(get1d)(bias, i));
+      }
+
+      /* do convolutions */
+      THTensor_(conv3Dmv)(outb, 1.0, 1.0, inb, weight, dT, dH, dW, "V", "X");
+    }
 
+    THTensor_(free)(inb);
+    THTensor_(free)(outb);
+  }
+  THTensor_(free)(outn);
+  
   return 1;
 }
 
 
 static int nn_(VolumetricConvolution_updateGradInput)(lua_State *L)
 {
+  THTensor *input = luaT_checkudata(L, 2, torch_Tensor);
   THTensor *gradOutput = luaT_checkudata(L, 3, torch_Tensor);  
   int dT = luaT_getfieldcheckint(L, 1, "dT");
   int dW = luaT_getfieldcheckint(L, 1, "dW");
@@ -58,12 +91,38 @@ static int nn_(VolumetricConvolution_updateGradInput)(lua_State *L)
   THTensor *weight = luaT_getfieldcheckudata(L, 1, "weight", torch_Tensor);
   THTensor *gradInput = luaT_getfieldcheckudata(L, 1, "gradInput", torch_Tensor);
   THTensor *tweight;
-  
-  THArgCheck( nOutputPlane == gradOutput->size[0], 1, "Number of output features is not equal to nOutputPlane" );
+
+  luaL_argcheck(L, gradOutput->nDimension == 4 || gradOutput->nDimension == 5, 
+		3, "4D or 5D (batch-mode) tensor expected");
+  int dimPlane = 0;
+  if (gradOutput->nDimension == 5) {
+    dimPlane++;
+  }
+  THArgCheck( nOutputPlane == gradOutput->size[dimPlane], 1, 
+	      "Number of output features is not equal to nOutputPlane" );
 
   /* gradient to input */
   tweight = THTensor_(newTranspose)(weight,0,1);
-  THTensor_(conv3Dmv)(gradInput, 0.0, 1.0, gradOutput, tweight, dT, dH, dW, "F", "C");
+  if (gradOutput->nDimension == 4) { /* non-batch mode */
+    THTensor_(conv3Dmv)(gradInput, 0.0, 1.0, gradOutput, tweight, dT, dH, dW, "F", "C");
+  } else { /* batch mode */
+    long nBatch = gradOutput->size[0];
+    THTensor *ginpb = THTensor_(new)();
+    THTensor *goutb = THTensor_(new)();
+    long j;
+    THTensor_(resize5d)(gradInput, input->size[0], input->size[1], input->size[2],
+			input->size[3], input->size[4]);
+
+#pragma omp parallel for private(j)
+    for (j=0; j<nBatch; j++) { /* loop over batches */
+      THTensor_(select)(ginpb,gradInput,0,j);
+      THTensor_(select)(goutb,gradOutput,0,j);
+      THTensor_(conv3Dmv)(ginpb, 0.0, 1.0, goutb, tweight, dT, dH, dW, "F", "C");
+    }
+    THTensor_(free)(ginpb);
+    THTensor_(free)(goutb);
+  }
+
   THTensor_(free)(tweight);
 
   return 1;
@@ -72,7 +131,7 @@ static int nn_(VolumetricConvolution_updateGradInput)(lua_State *L)
 static int nn_(VolumetricConvolution_accGradParameters)(lua_State *L)
 {
   THTensor *input = luaT_checkudata(L, 2, torch_Tensor);  
-  THTensor *gradOutput = luaT_checkudata(L, 3, torch_Tensor);  
+  THTensor *gradOutput = luaT_checkudata(L, 3, torch_Tensor);
   real scale = luaL_optnumber(L, 4, 1);
   int dT = luaT_getfieldcheckint(L, 1, "dT");
   int dW = luaT_getfieldcheckint(L, 1, "dW");
@@ -85,21 +144,55 @@ static int nn_(VolumetricConvolution_accGradParameters)(lua_State *L)
   long k;
   real *gradBias_data;
   THTensor* gradOutSlice;
-  
-  THArgCheck( nOutputPlane == gradOutput->size[0], 1, "Number of output features is not equal to nOutputPlane" );
-
-  /* gradient to bias */
-  gradBias_data = THTensor_(data)(gradBias);
-  gradOutSlice = THTensor_(new)();
-  for(k = 0; k < nOutputPlane; k++)
-  {
-    THTensor_(select)(gradOutSlice, gradOutput, 0, k);
-    gradBias_data[k] += scale*THTensor_(sumall)(gradOutSlice);
+  int dimPlane = 0;
+  if (gradOutput->nDimension == 5) {
+    dimPlane++;
   }
-  THTensor_(free)(gradOutSlice);
+  
+  THArgCheck( nOutputPlane == gradOutput->size[dimPlane], 1, 
+	      "Number of output features is not equal to nOutputPlane" );
 
-  /* gradient to kernels */
-  THTensor_(conv3DRevger)(gradWeight, 1.0, scale, input, gradOutput, dT, dH, dW);
+  
+  if (gradOutput->nDimension == 4) { /* non-batch mode */
+    /* gradient to bias */
+    gradBias_data = THTensor_(data)(gradBias);
+    gradOutSlice = THTensor_(new)();
+    for(k = 0; k < nOutputPlane; k++)
+      {
+	THTensor_(select)(gradOutSlice, gradOutput, 0, k);
+	gradBias_data[k] += scale*THTensor_(sumall)(gradOutSlice);
+      }
+    THTensor_(free)(gradOutSlice);
+    
+    /* gradient to kernels */
+    THTensor_(conv3DRevger)(gradWeight, 1.0, scale, input, gradOutput, dT, dH, dW);
+  } else { /* batch mode */
+    long nBatch = gradOutput->size[0];
+    THTensor *inpb = THTensor_(new)();
+    THTensor *goutb = THTensor_(new)();
+    long j;
+
+#pragma omp parallel for private(j)
+    for (j=0; j<nBatch; j++) { /* loop over batches */
+      THTensor_(select)(inpb,input,0,j);
+      THTensor_(select)(goutb,gradOutput,0,j);
+      
+      /* gradient to bias */
+      gradBias_data = THTensor_(data)(gradBias);
+      gradOutSlice = THTensor_(new)();
+      for(k = 0; k < nOutputPlane; k++)
+	{
+	  THTensor_(select)(gradOutSlice, goutb, 0, k);
+	  gradBias_data[k] += scale*THTensor_(sumall)(gradOutSlice);
+	}
+      THTensor_(free)(gradOutSlice);
+      
+      /* gradient to kernels */
+      THTensor_(conv3DRevger)(gradWeight, 1.0, scale, inpb, goutb, dT, dH, dW);
+    }
+    THTensor_(free)(inpb);
+    THTensor_(free)(goutb);
+  }
 
   return 0;
 }
diff --git a/test.lua b/test.lua
index a3c1c85..29de8bc 100644
--- a/test.lua
+++ b/test.lua
@@ -1253,9 +1253,9 @@ function nntest.SpatialFullConvolutionCompare()
 end
 
 local function batchcompare(smod, sin, plist)
-   local bs = torch.LongStorage(sin:size():size()+1)
+   local bs = torch.LongStorage(sin:dim()+1)
    bs[1] = 1
-   for i=1,sin:size():size() do bs[i+1] = sin:size()[i] end
+   for i=1,sin:dim() do bs[i+1] = sin:size()[i] end
    local bin = torch.Tensor(bs):copy(sin)
    local bmod = smod:clone()
 
@@ -1780,6 +1780,26 @@ function nntest.VolumetricConvolution()
    mytester:asserteq(0, berr, torch.typename(module) .. ' - i/o backward err ')
 end
 
+function nntest.VolumetricConvolutionBatchCompare()
+   local from = math.random(2,3)
+   local to = math.random(2,3)
+   local kt = math.random(3,4)
+   local ki = math.random(3,4)
+   local kj = math.random(3,4)
+   local st = math.random(2,3)
+   local si = math.random(2,3)
+   local sj = math.random(2,3)
+   local outt = math.random(3,4)
+   local outi = math.random(3,4)
+   local outj = math.random(3,4)
+   local int = (outt-1)*st+kt
+   local ini = (outi-1)*si+ki
+   local inj = (outj-1)*sj+kj
+   local module = nn.VolumetricConvolution(from, to, kt, ki, kj, st, si, sj)
+   local input = torch.randn(from, int, inj, ini)
+   batchcompare(module,input, {'weight','bias','gradWeight','gradBias'})   
+end
+
 function nntest.VolumetricMaxPooling()
    local from = math.random(2,3)
    local kt = math.random(3,4)