add batch information to getSpatialLocations

PyTorch/sparseconvnet/SCN/generic/Geometry/Metadata.cpp

@@ -63,9 +63,7 @@ extern "C" void scn_D_(setInputSpatialLocations)(void **m,
   assert((locations->size[1] == Dimension or
           locations->size[1] == 1 + Dimension) and
          "locations.size(0) must be either Dimension or Dimension+1");
-
   SCN_INITIALIZE_AND_REFERENCE(Metadata<Dimension>, m)
-
   Point<Dimension> p;
   auto &nActive = *_m.inputNActive;
   auto nPlanes = vecs->size[1];
@@ -73,7 +71,8 @@ extern "C" void scn_D_(setInputSpatialLocations)(void **m,
   auto v = THFloatTensor_data(vecs);
 
   if (locations->size[1] == Dimension) {
-    assert(_m.inputSG); // add points to current sample
+    // add points to current sample
+    assert(_m.inputSG);
     auto &mp = _m.inputSG->mp;
     for (uInt idx = 0; idx < locations->size[0]; ++idx) {
       for (int d = 0; d < Dimension; ++d)
@@ -101,9 +100,9 @@ extern "C" void scn_D_(setInputSpatialLocations)(void **m,
   }
 }
 
-extern "C" void scn_D_(getSpatialLocations)(void **m,
-                                            THLongTensor *spatialSize,
-                                            THLongTensor *locations) {
+extern "C" void scn_D_(getSpatialLocations)(void **m, THLongTensor *spatialSize,
+                                            THLongTensor *locations,
+                                            THLongTensor *batchIdxs) {
   SCN_INITIALIZE_AND_REFERENCE(Metadata<Dimension>, m)
   uInt nActive = _m.getNActive(spatialSize);
   auto &SGs = _m.getSparseGrid(spatialSize);
@@ -125,9 +124,9 @@ extern "C" void scn_D_(getSpatialLocations)(void **m,
   }
 }
 extern "C" void
-    scn_D_(createMetadataForDenseToSparse)(void **m, THLongTensor *spatialSize_,
-                                           THLongTensor *pad_,
-                                           THLongTensor *nz_, long batchSize) {
+scn_D_(createMetadataForDenseToSparse)(void **m, THLongTensor *spatialSize_,
+                                       THLongTensor *pad_, THLongTensor *nz_,
+                                       long batchSize) {
   SCN_INITIALIZE_AND_REFERENCE(Metadata<Dimension>, m)
   _m.setInputSpatialSize(spatialSize_);
   _m.inputSGs->resize(batchSize);

PyTorch/sparseconvnet/legacy/inputBatch.py

@@ -33,7 +33,25 @@ class InputBatch(SparseConvNetTensor):
             self.metadata.ffi, self.features, location, vector, overwrite)
 
     def setLocations(self, locations, vectors, overwrite=False):
-        l =locations.narrow(1,0,self.dimension)
+        """
+        To set n locations in d dimensions, locations can be
+        - A size (n,d) LongTensor, giving d-dimensional coordinates -- points
+          are added to the current sample, or
+        - A size (n,d+1) LongTensor; the extra column specifies the sample
+          number (within the minibatch of samples).
+
+          Example with d=3 and n=2:
+          Set
+          locations = LongTensor([[1,2,3],
+                                  [4,5,6]])
+          to add points to the current sample at (1,2,3) and (4,5,6).
+          Set
+          locations = LongTensor([[1,2,3,7],
+                                  [4,5,6,9]])
+          to add point (1,2,3) to sample 7, and (4,5,6) to sample 9 (0-indexed).
+
+        """
+        l = locations.narrow(1,0,self.dimension)
         assert l.min() >= 0 and (self.spatial_size.expand_as(l) - l).min() > 0
         dim_fn(self.dimension, 'setInputSpatialLocations')(
             self.metadata.ffi, self.features, locations, vectors, overwrite)

Torch/C.lua

@@ -25,35 +25,33 @@ return function (sparseconvnet)
   end
 
   cdef = [[
-  void scn_ptrCopyA(long *dst, void **src);
-  void scn_ptrCopyB(void **dst, long *src);
-  double scn_ruleBookBits();
-  void scn_2_drawCurve(void **m, THFloatTensor *features, THFloatTensor *stroke);
+void scn_ptrCopyA(long *dst, void **src);
+void scn_ptrCopyB(void **dst, long *src);
+double scn_ruleBookBits();
+void scn_2_drawCurve(void **m, THFloatTensor *features, THFloatTensor *stroke);
   ]]
   if fc then fc:write(cdef) end
   ffi.cdef(cdef)
   sparseconvnet.ruleBookBits=F['scn_ruleBookBits']()
 
   cdef = [[
-  double scn_DIMENSION_addSampleFromThresholdedTensor(
+double scn_DIMENSION_addSampleFromThresholdedTensor(
   void **m, THFloatTensor *features_, THFloatTensor *tensor_,
   THLongTensor *offset_, THLongTensor *spatialSize_, float threshold);
-  void scn_DIMENSION_batchAddSample(void **m);
-  void scn_DIMENSION_createMetadataForDenseToSparse(
+void scn_DIMENSION_batchAddSample(void **m);
+void scn_DIMENSION_createMetadataForDenseToSparse(
   void **m, THLongTensor *spatialSize_, THLongTensor *pad, THLongTensor *nz,
   long batchSize);
-  void scn_DIMENSION_freeMetadata(void **metadata);
-  void scn_DIMENSION_generateRuleBooks3s2(void **m);
-  void scn_DIMENSION_generateRuleBooks2s2(void **m);
-  void scn_DIMENSION_setInputSpatialSize(void **m, THLongTensor *spatialSize);
-  void scn_DIMENSION_setInputSpatialLocation(void **m,
-    THFloatTensor *features, THLongTensor *location, THFloatTensor *vec,
-    bool overwrite);
-  void scn_5_getSpatialLocations(void **m,
-    THLongTensor *spatialSize, THLongTensor *locations);
-  void scn_DIMENSION_setInputSpatialLocations(void **m,
-    THFloatTensor *features, THLongTensor *locations, THFloatTensor *vecs,
-    bool overwrite);
+void scn_DIMENSION_freeMetadata(void **metadata);
+void scn_DIMENSION_generateRuleBooks3s2(void **m);
+void scn_DIMENSION_generateRuleBooks2s2(void **m);
+void scn_DIMENSION_setInputSpatialSize(void **m, THLongTensor *spatialSize);
+void scn_DIMENSION_setInputSpatialLocation(void **m, THFloatTensor *features,
+  THLongTensor *location, THFloatTensor *vec, bool overwrite);
+void scn_DIMENSION_setInputSpatialLocations(void **m, THFloatTensor *features,
+  THLongTensor *locations, THFloatTensor *vecs, bool overwrite);
+void scn_DIMENSION_getSpatialLocations(void **m, THLongTensor *spatialSize,
+  THLongTensor *locations);
   ]]
 
   for DIMENSION = 1,10 do
@@ -66,60 +64,60 @@ return function (sparseconvnet)
   --type GPU half, float, double; int_cpu and int_gpu
 
   cdef = [[
-  void scn_ARCH_REAL_AffineReluTrivialConvolution_updateOutput(
+void scn_ARCH_REAL_AffineReluTrivialConvolution_updateOutput(
   THTensor *input_features, THTensor *output_features,
   THTensor *affineWeight, THTensor *affineBias, THTensor *convWeight);
-  void scn_ARCH_REAL_AffineReluTrivialConvolution_backward(
+void scn_ARCH_REAL_AffineReluTrivialConvolution_backward(
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, THTensor *affineWeight,
   THTensor *d_affineWeight, THTensor *affineBias, THTensor *d_affineBias,
   THTensor *convWeight, THTensor *d_convWeight, bool additiveGrad);
 
-  // BatchwiseMultiplicativeDropout
-  void scn_ARCH_REAL_BatchwiseMultiplicativeDropout_updateOutput(
+// BatchwiseMultiplicativeDropout
+void scn_ARCH_REAL_BatchwiseMultiplicativeDropout_updateOutput(
   THTensor *input_features, THTensor *output_features,
   THTensor *noise, long nPlanes, long input_stride, long output_stride,
   float alpha);
-  void scn_ARCH_REAL_BatchwiseMultiplicativeDropout_updateGradInput(
+void scn_ARCH_REAL_BatchwiseMultiplicativeDropout_updateGradInput(
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, THTensor *noise, long nPlanes,
   long input_stride, long output_stride, float alpha);
 
-  // BatchNormalization
-  void scn_ARCH_REAL_BatchNormalization_updateOutput(
+// BatchNormalization
+void scn_ARCH_REAL_BatchNormalization_updateOutput(
   THTensor *input_features, THTensor *output_features,
   THTensor *saveMean, THTensor *saveInvStd, THTensor *runningMean,
   THTensor *runningVar, THTensor *weight, THTensor *bias, REAL eps,
   REAL momentum, bool train, REAL leakiness);
-  void scn_ARCH_REAL_BatchNormalization_backward(
+void scn_ARCH_REAL_BatchNormalization_backward(
   THTensor *input_features, THTensor *d_input_features,
   THTensor *output_features, THTensor *d_output_features, THTensor *saveMean,
   THTensor *saveInvStd, THTensor *runningMean, THTensor *runningVar,
   THTensor *weight, THTensor *bias, THTensor *d_weight, THTensor *d_bias,
   REAL leakiness);
-  // BatchNormalizationInTensor
-  void scn_ARCH_REAL_BatchNormalizationInTensor_updateOutput(
+// BatchNormalizationInTensor
+void scn_ARCH_REAL_BatchNormalizationInTensor_updateOutput(
   THTensor *input_features, THTensor *output_features,
   THTensor *saveMean, THTensor *saveInvStd, THTensor *runningMean,
   THTensor *runningVar, THTensor *weight, THTensor *bias, REAL eps,
   REAL momentum, bool train, REAL leakiness);
 
-  // LeakyReLU
-  void scn_ARCH_REAL_LeakyReLU_updateOutput(
+// LeakyReLU
+void scn_ARCH_REAL_LeakyReLU_updateOutput(
   THTensor *input_features, THTensor *output_features, long n,
   float alpha);
-  void scn_ARCH_REAL_LeakyReLU_updateGradInput(
+void scn_ARCH_REAL_LeakyReLU_updateGradInput(
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, long n, float alpha);
 
-  // NetworkInNetwork
-  double scn_ARCH_REAL_NetworkInNetwork_updateOutput(
+// NetworkInNetwork
+double scn_ARCH_REAL_NetworkInNetwork_updateOutput(
   THTensor *input_features, THTensor *output_features,
   THTensor *weight, THTensor *bias);
-  void scn_ARCH_REAL_NetworkInNetwork_updateGradInput(
+void scn_ARCH_REAL_NetworkInNetwork_updateGradInput(
   THTensor *d_input_features, THTensor *d_output_features,
   THTensor *weight);
-  void scn_ARCH_REAL_NetworkInNetwork_accGradParameters(
+void scn_ARCH_REAL_NetworkInNetwork_accGradParameters(
   THTensor *input_features, THTensor *d_output_features,
   THTensor *d_weight, THTensor *d_bias);
   ]]
@@ -149,79 +147,79 @@ return function (sparseconvnet)
   end
 
   cdef = [[
-  // ActivePooling
-  void scn_ARCH_REAL_DIMENSIONActivePooling_updateOutput(
+// ActivePooling
+void scn_ARCH_REAL_DIMENSIONActivePooling_updateOutput(
   THLongTensor *inputSize, void **m, THTensor *input_features,
   THTensor *output_features, THITensor *rulesBuffer, bool average);
-  void scn_ARCH_REAL_DIMENSIONActivePooling_updateGradInput(
+void scn_ARCH_REAL_DIMENSIONActivePooling_updateGradInput(
   THLongTensor *inputSize, void **m,
   THTensor *d_input_features, THTensor *d_output_features,
   THITensor *rulesBuffer, bool average);
 
-  // Average Pooling
-  void scn_ARCH_REAL_DIMENSIONAveragePooling_updateOutput(
+// Average Pooling
+void scn_ARCH_REAL_DIMENSIONAveragePooling_updateOutput(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *poolSize, THLongTensor *poolStride, void **m,
   THTensor *input_features, THTensor *output_features, long nFeaturesToDrop,
   THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONAveragePooling_updateGradInput(
+void scn_ARCH_REAL_DIMENSIONAveragePooling_updateGradInput(
   THLongTensor * inputSize, THLongTensor * outputSize,
   THLongTensor * poolSize, THLongTensor * poolStride, void **m,
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, long nFeaturesToDrop,
   THITensor *rulesBuffer);
 
-  double scn_ARCH_REAL_DIMENSIONConvolution_updateOutput(
+double scn_ARCH_REAL_DIMENSIONConvolution_updateOutput(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *filterSize, THLongTensor *filterStride, void **m,
   THTensor *input_features, THTensor *output_features, THTensor *weight,
   THTensor *bias, long filterVolume, THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONConvolution_backward(
+void scn_ARCH_REAL_DIMENSIONConvolution_backward(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *filterSize, THLongTensor *filterStride, void **m,
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, THTensor *weight, THTensor *d_weight,
   THTensor *d_bias, long filterVolume, THITensor *rulesBuffer);
 
-  double scn_ARCH_REAL_DIMENSIONDeconvolution_updateOutput(
+double scn_ARCH_REAL_DIMENSIONDeconvolution_updateOutput(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *filterSize, THLongTensor *filterStride, void **m,
   THTensor *input_features, THTensor *output_features, THTensor *weight,
   THTensor *bias, long filterVolume, THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONDeconvolution_backward(
+void scn_ARCH_REAL_DIMENSIONDeconvolution_backward(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *filterSize, THLongTensor *filterStride, void **m,
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, THTensor *weight, THTensor *d_weight,
   THTensor *d_bias, long filterVolume, THITensor *rulesBuffer);
 
-  // Max Pooling
-  void scn_ARCH_REAL_DIMENSIONMaxPooling_updateOutput(
+// Max Pooling
+void scn_ARCH_REAL_DIMENSIONMaxPooling_updateOutput(
   THLongTensor *inputSize, THLongTensor *outputSize,
   THLongTensor *poolSize, THLongTensor *poolStride, void **m,
   THTensor *input_features, THTensor *output_features, long nFeaturesToDrop,
   THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONMaxPooling_updateGradInput(
+void scn_ARCH_REAL_DIMENSIONMaxPooling_updateGradInput(
   THLongTensor * inputSize, THLongTensor * outputSize,
   THLongTensor * poolSize, THLongTensor * poolStride, void **m,
   THTensor *input_features, THTensor *d_input_features,
   THTensor *output_features, THTensor *d_output_features,
   long nFeaturesToDrop, THITensor *rulesBuffer);
 
-  // SparseToDense
-  void scn_ARCH_REAL_DIMENSIONSparseToDense_updateOutput(
+// SparseToDense
+void scn_ARCH_REAL_DIMENSIONSparseToDense_updateOutput(
   THLongTensor *inputSize, void **m, THTensor *input_features,
   THTensor *output_features, THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONSparseToDense_updateGradInput(
+void scn_ARCH_REAL_DIMENSIONSparseToDense_updateGradInput(
   THLongTensor *inputSize, void **m, THTensor *input_features,
   THTensor *d_input_features, THTensor *d_output_features,
   THITensor *rulesBuffer);
 
-  double scn_ARCH_REAL_DIMENSIONValidConvolution_updateOutput(
+double scn_ARCH_REAL_DIMENSIONValidConvolution_updateOutput(
   THLongTensor *inputSize, THLongTensor *filterSize, void **m,
   THTensor *input_features, THTensor *output_features, THTensor *weight,
   THTensor *bias, long filterVolume, THITensor *rulesBuffer);
-  void scn_ARCH_REAL_DIMENSIONValidConvolution_backward(
+void scn_ARCH_REAL_DIMENSIONValidConvolution_backward(
   THLongTensor *inputSize, THLongTensor *filterSize, void **m,
   THTensor *input_features, THTensor *d_input_features,
   THTensor *d_output_features, THTensor *weight, THTensor *d_weight,