fix some bug in python. add USAGE.md

README.md

@@ -84,6 +84,10 @@ You need to rebuild ```cumm``` first if you are build along a CUDA version that
 4. run ```$Env:SPCONV_DISABLE_JIT = "1"```
 5. run ```python setup.py install```/```pip install -e .```/```python setup.py bdist_wheel```+```pip install dists/xxx.whl```
 
+## Documents
+
+see docs/USAGE.md.
+
 ## Note
 
 The work is done when the author is an employee at Tusimple.

docs/USAGE.md

+<!--
+ Copyright 2021 Yan Yan
+ 
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+ 
+     http://www.apache.org/licenses/LICENSE-2.0
+ 
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+-->
+
+# Usage
+
+## Concept
+
+* Sparse Conv Tensor: like hybird [torch.sparse_coo_tensor](https://pytorch.org/docs/stable/sparse.html#sparse-coo-docs) but only have two difference: 1. SparseConvTensor only have one dense dim, 2. indice of SparseConvTensor is transposed. see torch doc for more details.
+
+* Sparse Convolution: equivalent to perform dense convolution when you convert SparseConvTensor to dense. Sparse Convolution only run calculation on valid data.
+
+* Submanifold Convolution (SubMConv): like Sparse Convolution but indices keeps same. imagine that you copy same spatial structure to output, then iterate them, get input coordinates by conv rule, finally apply convolution **ONLY** in these output coordinates.
+
+
+
+## SparseConvTensor
+
+* features: ```[N, num_channels]``` tensor.
+
+* indices: ```[N, (batch_idx + x + y + z)]``` coordinate tensor with batch axis. note that the coordinates xyz order MUST match spatial shape and conv params such as kernel size
+
+```Python
+features = # your features with shape [N, num_channels]
+indices = # your indices/coordinates with shape [N, ndim + 1], batch index must be put in indices[:, 0]
+spatial_shape = # spatial shape of your sparse tensor, spatial_shape[i] is shape of indices[:, 1 + i].
+batch_size = # batch size of your sparse tensor.
+x = spconv.SparseConvTensor(features, indices, spatial_shape, batch_size)
+x_dense_NCHW = x.dense() # convert sparse tensor to dense NCHW tensor.
+```
+
+
+### Sparse Convolution
+
+```Python
+import spconv
+from torch import nn
+class ExampleNet(nn.Module):
+    def __init__(self, shape):
+        super().__init__()
+        self.net = spconv.SparseSequential(
+            spconv.SparseConv3d(32, 64, 3), # just like nn.Conv3d but don't support group and all([d > 1, s > 1])
+            nn.BatchNorm1d(64), # non-spatial layers can be used directly in SparseSequential.
+            nn.ReLU(),
+            spconv.SubMConv3d(64, 64, 3, indice_key="subm0"),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            # when use submanifold convolutions, their indices can be shared to save indices generation time.
+            spconv.SubMConv3d(64, 64, 3, indice_key="subm0"),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            spconv.SparseConvTranspose3d(64, 64, 3, 2),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            spconv.ToDense(), # convert spconv tensor to dense and convert it to NCHW format.
+            nn.Conv3d(64, 64, 3),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+        )
+        self.shape = shape
+
+    def forward(self, features, coors, batch_size):
+        coors = coors.int() # unlike torch, this library only accept int coordinates.
+        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
+        return self.net(x)# .dense()
+```
+
+### Inverse Convolution
+
+Inverse sparse convolution means "inv" of sparse convolution. the output of inverse convolution contains same indices as input of sparse convolution.
+
+Inverse convolution usually used in semantic segmentation.
+
+```Python
+class ExampleNet(nn.Module):
+    def __init__(self, shape):
+        super().__init__()
+        self.net = spconv.SparseSequential(
+            spconv.SparseConv3d(32, 64, 3, 2, indice_key="cp0"),
+            spconv.SparseInverseConv3d(64, 32, 3, indice_key="cp0"), # need provide kernel size to create weight
+        )
+        self.shape = shape
+
+    def forward(self, features, coors, batch_size):
+        coors = coors.int()
+        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
+        return self.net(x)
+```
+
+
+### Utility functions
+
+* convert point cloud to voxel
+
+voxel generator in spconv generate indices in **ZYX** order, the params format are **XYZ**.
+
+voxel generator in spconv takes a ```tv.Tensor``` return a ```tv.Tensor```, this tensor reference to a **permanent** storage in generator.
+
+
+```Python
+from spconv.utils import Point2VoxelCPU3d
+# this generator generate ZYX indices.
+gen = Point2VoxelCPU3d(
+    vsize_xyz=[0.1, 0.1, 0.1], 
+    coors_range_xyz=[-80, -80, -2, 80, 80, 6], 
+    num_point_features=3, 
+    max_num_voxels=5000, 
+    max_num_points_per_voxel=5)
+pc = np.random.uniform(-10, 10, size=[1000, 3])
+pc_tv = tv.from_numpy(pc)
+voxels, coords, num_points_per_voxel = gen.generate(pc_tv)
+
+# get numpy
+voxels_np = voxels.numpy_view() # no copy, but become invalid if generator is destroyed.
+voxels_np = voxels.numpy() # will perform copy
+```

spconv/pytorch/conv.py

@@ -92,10 +92,6 @@ class SparseConvolution(SparseModule):
             dilation = [dilation] * ndim
         if not isinstance(output_padding, (list, tuple)):
             output_padding = [output_padding] * ndim
-
-        for d, s in zip(dilation, stride):
-            assert any([s == 1, d == 1]), "don't support this."
-
         self.ndim = ndim
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -186,11 +182,11 @@ class SparseConvolution(SparseModule):
             else:
                 features = torch.mm(
                     input.features,
-                    self.weight.view(self.in_channels, self.out_channels).T)
+                    self.weight.view(self.in_channels, self.out_channels))
 
             if self.bias is not None:
                 features += self.bias
-            out_tensor.features = features
+            out_tensor = out_tensor.replace_feature(features)
             return out_tensor
         datas = input.find_indice_pair(self.indice_key)
         if self.inverse:
@@ -272,8 +268,7 @@ class SparseConvolution(SparseModule):
                 features.shape[0])
             out_tensor.benchmark_record[self.name]["num_out_points"].append(
                 out_features.shape[0])
-
-        out_tensor.features = out_features
+        out_tensor = out_tensor.replace_feature(out_features)
         out_tensor.indices = outids
         out_tensor.spatial_shape = out_spatial_shape
         return out_tensor

spconv/pytorch/core_fx.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 import torch
@@ -47,13 +47,14 @@ def scatter_nd(indices, updates, shape):
 
 class SparseConvTensor(metaclass=ProxyableClassMeta):
     def __init__(self,
-                 features,
-                 indices,
-                 spatial_shape,
-                 batch_size,
-                 grid=None,
-                 voxel_num=None,
-                 benchmark=False):
+                 features: torch.Tensor,
+                 indices: torch.Tensor,
+                 spatial_shape: List[int],
+                 batch_size: int,
+                 grid: Optional[torch.Tensor]=None,
+                 voxel_num: Optional[torch.Tensor]=None,
+                 indice_dict: Optional[dict] = None,
+                 benchmark: bool=False):
         """
         Args:
             features: [num_points, num_features] feature tensor
@@ -69,7 +70,9 @@ class SparseConvTensor(metaclass=ProxyableClassMeta):
         self.indices = indices
         self.spatial_shape = spatial_shape
         self.batch_size = batch_size
-        self.indice_dict = {}
+        if indice_dict is None:
+            indice_dict = {}
+        self.indice_dict = indice_dict
         if grid is None:
             grid = torch.Tensor()  # empty tensor
         self.grid = grid
@@ -101,11 +104,11 @@ class SparseConvTensor(metaclass=ProxyableClassMeta):
         """create sparse tensor fron channel last dense tensor by to_sparse
         x must be NHWC tensor, channel last
         """
-        x = x.to_sparse(x.ndim - 1)
-        spatial_shape = x.shape[1:-1]
-        batch_size = x.shape[0]
-        indices_th = x.indices().permute(1, 0).contiguous().int()
-        features_th = x.values()
+        x_sp = x.to_sparse(x.ndim - 1)
+        spatial_shape = list(x_sp.shape[1:-1])
+        batch_size = x_sp.shape[0]
+        indices_th = x_sp.indices().permute(1, 0).contiguous().int()
+        features_th = x_sp.values()
         return cls(features_th, indices_th, spatial_shape, batch_size)
 
     @property
@@ -119,7 +122,7 @@ class SparseConvTensor(metaclass=ProxyableClassMeta):
             return self.indice_dict[key]
         return None
 
-    def dense(self, channels_first=True):
+    def dense(self, channels_first: bool=True):
         output_shape = [self.batch_size] + list(
             self.spatial_shape) + [self.features.shape[1]]
         res = scatter_nd(
@@ -142,8 +145,6 @@ class SparseConvTensor(metaclass=ProxyableClassMeta):
         """create a new spconv tensor with all member unchanged"""
         tensor = SparseConvTensor(self.features, self.indices,
                                   self.spatial_shape, self.batch_size,
-                                  self.grid, self.benchmark)
+                                  self.grid, self.voxel_num, self.indice_dict, self.benchmark)
         tensor.benchmark_record = self.benchmark_record
-        tensor.indice_dict = self.indice_dict
-        tensor.voxel_num = self.voxel_num
         return tensor

spconv/pytorch/core_nofx.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 import torch
@@ -44,15 +44,16 @@ def scatter_nd(indices, updates, shape):
     return ret
 
 
-class SparseConvTensor(object):
+class SparseConvTensor:
     def __init__(self,
-                 features,
-                 indices,
-                 spatial_shape,
-                 batch_size,
-                 grid=None,
-                 voxel_num=None,
-                 benchmark=False):
+                 features: torch.Tensor,
+                 indices: torch.Tensor,
+                 spatial_shape: List[int],
+                 batch_size: int,
+                 grid: Optional[torch.Tensor]=None,
+                 voxel_num: Optional[torch.Tensor]=None,
+                 indice_dict: Optional[dict] = None,
+                 benchmark: bool=False):
         """
         Args:
             features: [num_points, num_features] feature tensor
@@ -68,16 +69,18 @@ class SparseConvTensor(object):
         self.indices = indices
         self.spatial_shape = spatial_shape
         self.batch_size = batch_size
-        self.indice_dict = {}
+        if indice_dict is None:
+            indice_dict = {}
+        self.indice_dict = indice_dict
         if grid is None:
             grid = torch.Tensor()  # empty tensor
         self.grid = grid
-        self.voxel_num = voxel_num
+        self.voxel_num = voxel_num # for tensorrt
         self.benchmark = benchmark
         self.benchmark_record = {}
 
     def replace_feature(self, feature):
-        """we need to replace x.features = F.relu(x) with x = x.replace_feature(F.relu(x))
+        """we need to replace x.features = F.relu(x) with x = x.replace_feature(F.relu(x.features))
         due to limit of torch.fx
         """
         new_spt = SparseConvTensor(feature, self.indices, self.spatial_shape, self.batch_size, self.grid, self.voxel_num, self.indice_dict)
@@ -91,7 +94,7 @@ class SparseConvTensor(object):
 
     @features.setter
     def features(self, val):
-        msg = ("you can't set feature directly, use 'x = x.replace_feature(F.relu(x.feature))'"
+        msg = ("you can't set feature directly, use 'x = x.replace_feature(your_new_feature)'"
                 " to generate new SparseConvTensor instead.")
         raise ValueError(msg)
 
@@ -100,11 +103,11 @@ class SparseConvTensor(object):
         """create sparse tensor fron channel last dense tensor by to_sparse
         x must be NHWC tensor, channel last
         """
-        x = x.to_sparse(x.ndim - 1)
-        spatial_shape = x.shape[1:-1]
-        batch_size = x.shape[0]
-        indices_th = x.indices().permute(1, 0).contiguous().int()
-        features_th = x.values()
+        x_sp = x.to_sparse(x.ndim - 1)
+        spatial_shape = list(x_sp.shape[1:-1])
+        batch_size = x_sp.shape[0]
+        indices_th = x_sp.indices().permute(1, 0).contiguous().int()
+        features_th = x_sp.values()
         return cls(features_th, indices_th, spatial_shape, batch_size)
 
     @property
@@ -118,7 +121,7 @@ class SparseConvTensor(object):
             return self.indice_dict[key]
         return None
 
-    def dense(self, channels_first=True):
+    def dense(self, channels_first: bool=True):
         output_shape = [self.batch_size] + list(
             self.spatial_shape) + [self.features.shape[1]]
         res = scatter_nd(
@@ -131,6 +134,7 @@ class SparseConvTensor(object):
         trans_params.insert(1, ndim + 1)
         return res.permute(*trans_params).contiguous()
 
+    # remove this due to limit of torch.fx
     # @property
     # def sparity(self):
     #     return self.indices.shape[0] / np.prod(
@@ -140,7 +144,6 @@ class SparseConvTensor(object):
         """create a new spconv tensor with all member unchanged"""
         tensor = SparseConvTensor(self.features, self.indices,
                                   self.spatial_shape, self.batch_size,
-                                  self.grid, self.benchmark)
+                                  self.grid, self.voxel_num, self.indice_dict, self.benchmark)
         tensor.benchmark_record = self.benchmark_record
-        tensor.indice_dict = self.indice_dict
         return tensor

spconv/pytorch/modules.py

@@ -100,7 +100,7 @@ class SparseSequential(SparseModule):
             if name in self._modules:
                 raise ValueError("name exists.")
             self.add_module(name, module)
-        self._sparity_dict = {}
+        # self._sparity_dict = {}
 
     def __getitem__(self, idx):
         if not (-len(self) <= idx < len(self)):
@@ -115,9 +115,9 @@ class SparseSequential(SparseModule):
     def __len__(self):
         return len(self._modules)
 
-    @property
-    def sparity_dict(self):
-        return self._sparity_dict
+    # @property
+    # def sparity_dict(self):
+    #     return self._sparity_dict
 
     def add(self, module, name=None):
         if name is None:
@@ -133,12 +133,12 @@ class SparseSequential(SparseModule):
                     input = module(input)
                 else:
                     assert isinstance(input, spconv.SparseConvTensor)
-                    self._sparity_dict[k] = input.sparity
+                    # self._sparity_dict[k] = input.sparity
                     input = module(input)
             else:
                 if isinstance(input, spconv.SparseConvTensor):
                     if input.indices.shape[0] != 0:
-                        input.features = module(input.features)
+                        input.replace_feature(module(input.features))
                 else:
                     input = module(input)
         return input

spconv/pytorch/pool.py

@@ -135,8 +135,7 @@ class SparseMaxPool(SparseModule):
                 features.shape[0])
             out_tensor.benchmark_record[self.name]["num_out_points"].append(
                 out_features.shape[0])
-
-        out_tensor.features = out_features
+        out_tensor = out_tensor.replace_feature(out_features)
         out_tensor.indices = outids
         out_tensor.spatial_shape = out_spatial_shape
         return out_tensor

test/benchmark.py

@@ -151,6 +151,9 @@ class Net(nn.Module):
                               bias=False,
                               indice_key="c6",
                               algo=algo),
+            spconv.SparseInverseConv3d(256, 128, 3, indice_key="c6", bias=False),
+            spconv.SparseInverseConv3d(128, 64, 3, indice_key="c5", bias=False),
+
         )
         max_batch_size = 1
         # grid (dense map) is used for indice generation. use pre-allocated grid can run faster.