add ConcaTable, JoinTable, Identity functions

spconv/__init__.py

@@ -23,6 +23,8 @@ from spconv.conv import SparseConvTranspose2d, SparseConvTranspose3d
 from spconv.conv import SparseInverseConv2d, SparseInverseConv3d
 from spconv.modules import SparseModule, SparseSequential
 from spconv.pool import SparseMaxPool2d, SparseMaxPool3d
+from spconv.tables import ConcatTable, JoinTable
+from spconv.identity import Identity
 
 from spconv import ops
 

spconv/identity.py

+# Copyright 2016-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch.nn import Module
+
+
+class Identity(Module):
+    def forward(self, input):
+        return input
+
+    def input_spatial_size(self, out_size):
+        return out_size

spconv/tables.py

+
+from torch.autograd import Function
+#from torch.nn import Module
+from spconv.modules import SparseModule
+import spconv
+import torch
+
+class JoinTable(SparseModule):# Module):
+    def forward(self, input):
+        output = spconv.SparseConvTensor(
+                torch.cat([i.features for i in input],1), input[1].indices,
+                input[1].spatial_shape, input[0].batch_size )
+        output.indice_dict = input[1].indice_dict
+        output.grid = input[1].grid
+        return output
+
+    def input_spatial_size(self, out_size):
+        return out_size
+
+
+class AddTable(SparseModule): # Module):
+    def forward(self, input):
+        output = spconv.SparseConvTensor(
+                sum([i.features for i in input]), input[1].indices,
+                input[1].spatial_shape, input[1].batch_size )
+        output.indice_dict = input[1].indice_dict
+        output.grid = input[1].grid
+
+        return output
+
+    def input_spatial_size(self, out_size):
+        return out_size
+
+
+class ConcatTable(SparseModule): # Module):
+    def forward(self, input):
+        return [module(input) for module in self._modules.values()]
+
+    def add(self, module):
+        self._modules[str(len(self._modules))] = module
+        return self
+
+    def input_spatial_size(self, out_size):
+        return self._modules['0'].input_spatial_size(out_size)