SparseConv from Open3D (#29)

Author: dkurt

.github/workflows/main.yml

@@ -12,7 +12,7 @@ on:
 
 env:
   OPENVINO_VERSION: 2021.4.2
-  VERSION: 2021.4.2.3
+  VERSION: 2021.4.2.4
   DIST_VERSION: 2021.4.752
   DIST_WIN: https://registrationcenter-download.intel.com/akdlm/irc_nas/18320/w_openvino_toolkit_p_2021.4.752.exe
   DIST_MAC: https://registrationcenter-download.intel.com/akdlm/irc_nas/18317/m_openvino_toolkit_p_2021.4.752.dmg
@@ -134,9 +134,6 @@ jobs:
   test_lnx:
     needs: build_lnx
     runs-on: ubuntu-18.04
-    strategy:
-      matrix:
-        torch-version: [1.7.0, 1.10.0]
 
     steps:
     - uses: actions/checkout@v2
@@ -149,7 +146,7 @@ jobs:
       run: |
         sudo apt-get install -y python3-setuptools libopencv-dev
         python3 -m pip install --upgrade pip
-        python3 -m pip install torch==${{ matrix.torch-version }} torchvision
+        python3 -m pip install -r tests/requirements.txt
         python3 -m pip install -U protobuf
         python3 -m pip install openvino-dev[onnx]==${{env.OPENVINO_VERSION}}
 
@@ -161,7 +158,7 @@ jobs:
 
     - name: Test
       run: |
-        python3 -m unittest run_tests.py
+        python3 -m pytest tests/run_tests.py
 
   publish:
     if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
@@ -177,4 +174,4 @@ jobs:
         run: |
           python3 -m pip install --upgrade pip
           python3 -m pip install twine
-          python3 -m twine upload wheel*/*.whl --skip-existing
+          python3 -m twine upload --repository testpypi wheel*/*.whl --skip-existing

README.md

@@ -6,6 +6,7 @@ Repository with guides to enable some layers from PyTorch in Intel OpenVINO:
 * [torch.fft](examples/fft)
 * [nn.functional.grid_sample](https://github.com/dkurt/openvino_pytorch_layers/tree/master/examples/grid_sample)
 * [torchvision.ops.DeformConv2d](examples/deformable_conv)
+* [SparseConv](examples/sparse_conv) from [Open3D](https://github.com/isl-org/Open3D)
 
 
 ## OpenVINO Model Optimizer extension

examples/calculate_grid/calculate_grid.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CalculateGrid(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, in_positions):
+        return g.op("org.open3d::calculate_grid", in_positions)
+
+    @staticmethod
+    def forward(self, in_positions):
+        filter = torch.Tensor([[-1, -1, -1], [-1, -1, 0], [-1, 0, -1], [-1, 0, 0],
+                               [0, -1, -1], [0, -1, 0], [0, 0, -1],
+                               [0, 0, 0]]).to(in_positions.device)
+
+        out_pos = in_positions.long().repeat(1, filter.shape[0]).reshape(-1, 3)
+        filter = filter.repeat(in_positions.shape[0], 1)
+
+        out_pos = out_pos + filter
+        out_pos = out_pos[out_pos.min(1).values >= 0]
+        out_pos = out_pos[(~((out_pos.long() % 2).bool()).any(1))]
+        out_pos = torch.unique(out_pos, dim=0)
+
+        return out_pos + 0.5

examples/calculate_grid/export_model.py

@@ -0,0 +1,40 @@
+import numpy as np
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .calculate_grid import CalculateGrid
+
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.calculate_grid = CalculateGrid()
+
+    def forward(self, x):
+        return self.calculate_grid.apply(x)
+
+
+def export(num_points, max_grid_extent):
+    # Generate a list of unique positions and add a mantissa
+    np.random.seed(32)
+    torch.manual_seed(11)
+
+    inp_pos = np.random.randint(0, max_grid_extent, [num_points, 3])
+    inp_pos = torch.tensor(inp_pos) + torch.rand(inp_pos.shape, dtype=torch.float32) # [0, 1)
+
+    model = MyModel()
+    with torch.no_grad():
+        torch.onnx.export(model, (inp_pos), 'model.onnx',
+                          input_names=['input'],
+                          output_names=['output'],
+                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+
+    ref = model(inp_pos).detach().numpy()
+
+    # Pad values with espetial end line (-1, 0, 0) and zeros
+    ref = np.concatenate((ref, [[-1, 0, 0]]))
+    ref = np.pad(ref, ((0, inp_pos.shape[0] - ref.shape[0]), (0, 0)))
+
+    np.save('inp', inp_pos.detach().numpy())
+    np.save('ref', ref)

examples/sparse_conv/export_model.py

@@ -0,0 +1,69 @@
+import numpy as np
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .sparse_conv import SparseConvONNX, SparseConvTransposeONNX
+
+
+def export(num_inp_points, num_out_points, max_grid_extent, in_channels,
+           filters, kernel_size, normalize, transpose):
+    np.random.seed(324)
+    torch.manual_seed(32)
+
+    if transpose:
+        sparse_conv = SparseConvTransposeONNX(in_channels=in_channels,
+                                              filters=filters,
+                                              kernel_size=kernel_size,
+                                              use_bias=False,
+                                              normalize=False)
+    else:
+        sparse_conv = SparseConvONNX(in_channels=in_channels,
+                                     filters=filters,
+                                     kernel_size=kernel_size,
+                                     use_bias=False,
+                                     normalize=False)
+
+    # Generate a list of unique positions and add a mantissa
+    def gen_pos(num_points):
+        inp_pos = np.random.randint(0, max_grid_extent, [num_points, 3])
+        inp_pos = np.unique(inp_pos, axis=0).astype(np.float32)
+        inp_pos = torch.tensor(inp_pos) + torch.rand(inp_pos.shape, dtype=torch.float32) # [0, 1)
+        return inp_pos
+
+    inp_pos = gen_pos(num_inp_points)
+    out_pos = gen_pos(num_out_points) if num_out_points else inp_pos
+
+    features = torch.randn([inp_pos.shape[0], in_channels])
+
+    voxel_size = torch.tensor(1.0)
+    sparse_conv.eval()
+
+    new_kernel = torch.randn(sparse_conv.state_dict()["kernel"].shape)
+    sparse_conv.load_state_dict({"kernel": new_kernel,
+                                 "offset": sparse_conv.state_dict()["offset"]})
+
+    with torch.no_grad():
+        torch.onnx.export(sparse_conv, (features, inp_pos, out_pos, voxel_size), 'model.onnx',
+                          input_names=['input', 'input1', 'input2', 'voxel_size'],
+                          output_names=['output'],
+                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+
+    ref = sparse_conv(features, inp_pos, out_pos, voxel_size)
+    np.save('inp', features.detach().numpy())
+    np.save('inp1', inp_pos.detach().numpy())
+    np.save('inp2', out_pos.detach().numpy())
+    np.save('ref', ref.detach().numpy())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Generate ONNX model and test data')
+    parser.add_argument('--num_points', type=int)
+    parser.add_argument('--max_grid_extent', type=int)
+    parser.add_argument('--in_channels', type=int)
+    parser.add_argument('--filters', type=int)
+    parser.add_argument('--kernel_size', type=int)
+    args = parser.parse_args()
+
+    export(args.num_points, args.max_grid_extent,
+           args.in_channels, args.filters, args.kernel_size)

examples/sparse_conv/sparse_conv.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from open3d.ml.torch.layers import SparseConv, SparseConvTranspose
+
+class SparseConvFunc(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, cls, feat, in_pos, out_pos, voxel_size):
+        kernel = cls.state_dict()["kernel"]
+        offset = cls.state_dict()["offset"]
+        kernel = g.op("Constant", value_t=kernel)
+        offset = g.op("Constant", value_t=offset)
+        return g.op("org.open3d::SparseConv", feat, in_pos, out_pos, kernel, offset)
+
+    @staticmethod
+    def forward(self, cls, feat, in_pos, out_pos, voxel_size):
+        return cls.origin_forward(feat, in_pos, out_pos, voxel_size)
+
+
+class SparseConvONNX(SparseConv):
+    """
+    This is a support class which helps export network with SparseConv in ONNX format.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.origin_forward = super().forward
+
+    def forward(self, feat, in_pos, out_pos, voxel_size):
+        return SparseConvFunc.apply(self, feat, in_pos, out_pos, voxel_size)
+
+
+class SparseConvTransposeFunc(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, cls, feat, in_pos, out_pos, voxel_size):
+        kernel = cls.state_dict()["kernel"]
+        offset = cls.state_dict()["offset"]
+        kernel = g.op("Constant", value_t=kernel)
+        offset = g.op("Constant", value_t=offset)
+        return g.op("org.open3d::SparseConvTranspose", feat, in_pos, out_pos, kernel, offset)
+
+    @staticmethod
+    def forward(self, cls, feat, in_pos, out_pos, voxel_size):
+        return cls.origin_forward(feat, in_pos, out_pos, voxel_size)
+
+
+class SparseConvTransposeONNX(SparseConvTranspose):
+    """
+    This is a support class which helps export network with SparseConvTranspose in ONNX format.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.origin_forward = super().forward
+
+    def forward(self, feat, in_pos, out_pos, voxel_size):
+        return SparseConvTransposeFunc.apply(self, feat, in_pos, out_pos, voxel_size)

run_tests.py

@@ -0,0 +1,5 @@
+torch==1.8.1
+torchvision==0.9.1
+open3d==0.14.1
+tensorboard
+pytest