Merge pull request #658 from acivgin1/master for spconv 2.0

Add support for spconv 2.0

pcdet/datasets/dataset.py

@@ -31,7 +31,8 @@ class DatasetTemplate(torch_data.Dataset):
             self.root_path, self.dataset_cfg.DATA_AUGMENTOR, self.class_names, logger=self.logger
         ) if self.training else None
         self.data_processor = DataProcessor(
-            self.dataset_cfg.DATA_PROCESSOR, point_cloud_range=self.point_cloud_range, training=self.training
+            self.dataset_cfg.DATA_PROCESSOR, point_cloud_range=self.point_cloud_range,
+            training=self.training, num_point_features=self.point_feature_encoder.num_point_features
         )
 
         self.grid_size = self.data_processor.grid_size

pcdet/datasets/processor/data_processor.py

@@ -5,14 +5,72 @@ from skimage import transform
 
 from ...utils import box_utils, common_utils
 
+tv = None
+try:
+    import cumm.tensorview as tv
+except:
+    pass
+
+
+class VoxelGeneratorWrapper():
+    def __init__(self, vsize_xyz, coors_range_xyz, num_point_features, max_num_points_per_voxel, max_num_voxels):
+        try:
+            from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
+            self.spconv_ver = 1
+        except:
+            try:
+                from spconv.utils import VoxelGenerator
+                self.spconv_ver = 1
+            except:
+                from spconv.utils import Point2VoxelCPU3d as VoxelGenerator
+                self.spconv_ver = 2
+
+        if self.spconv_ver == 1:
+            self._voxel_generator = VoxelGenerator(
+                voxel_size=vsize_xyz,
+                point_cloud_range=coors_range_xyz,
+                max_num_points=max_num_points_per_voxel,
+                max_voxels=max_num_voxels
+            )
+        else:
+            self._voxel_generator = VoxelGenerator(
+                vsize_xyz=vsize_xyz,
+                coors_range_xyz=coors_range_xyz,
+                num_point_features=num_point_features,
+                max_num_points_per_voxel=max_num_points_per_voxel,
+                max_num_voxels=max_num_voxels
+            )
+
+    def generate(self, points):
+        if self.spconv_ver == 1:
+            voxel_output = self._voxel_generator.generate(points)
+            if isinstance(voxel_output, dict):
+                voxels, coordinates, num_points = \
+                    voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
+            else:
+                voxels, coordinates, num_points = voxel_output
+        else:
+            assert tv is not None, f"Unexpected error, library: 'cumm' wasn't imported properly."
+            voxel_output = self._voxel_generator.point_to_voxel(tv.from_numpy(points))
+            tv_voxels, tv_coordinates, tv_num_points = voxel_output
+            # make copy with numpy(), since numpy_view() will disappear as soon as the generator is deleted
+            voxels = tv_voxels.numpy()
+            coordinates = tv_coordinates.numpy()
+            num_points = tv_num_points.numpy()
+        return voxels, coordinates, num_points
+
 
 class DataProcessor(object):
-    def __init__(self, processor_configs, point_cloud_range, training):
+    def __init__(self, processor_configs, point_cloud_range, training, num_point_features):
         self.point_cloud_range = point_cloud_range
         self.training = training
+        self.num_point_features = num_point_features
         self.mode = 'train' if training else 'test'
         self.grid_size = self.voxel_size = None
         self.data_processor_queue = []
+
+        self.voxel_generator = None
+
         for cur_cfg in processor_configs:
             cur_processor = getattr(self, cur_cfg.NAME)(config=cur_cfg)
             self.data_processor_queue.append(cur_processor)
@@ -44,31 +102,27 @@ class DataProcessor(object):
 
         return data_dict
 
-    def transform_points_to_voxels(self, data_dict=None, config=None, voxel_generator=None):
+    def transform_points_to_voxels(self, data_dict=None, config=None):
         if data_dict is None:
-            try:
-                from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
-            except:
-                from spconv.utils import VoxelGenerator
-
-            voxel_generator = VoxelGenerator(
-                voxel_size=config.VOXEL_SIZE,
-                point_cloud_range=self.point_cloud_range,
-                max_num_points=config.MAX_POINTS_PER_VOXEL,
-                max_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode]
-            )
             grid_size = (self.point_cloud_range[3:6] - self.point_cloud_range[0:3]) / np.array(config.VOXEL_SIZE)
             self.grid_size = np.round(grid_size).astype(np.int64)
             self.voxel_size = config.VOXEL_SIZE
-            return partial(self.transform_points_to_voxels, voxel_generator=voxel_generator)
+            # just bind the config, we will create the VoxelGeneratorWrapper later,
+            # to avoid pickling issues in multiprocess spawn
+            return partial(self.transform_points_to_voxels, config=config)
+
+        if self.voxel_generator is None:
+            self.voxel_generator = VoxelGeneratorWrapper(
+                vsize_xyz=config.VOXEL_SIZE,
+                coors_range_xyz=self.point_cloud_range,
+                num_point_features=self.num_point_features,
+                max_num_points_per_voxel=config.MAX_POINTS_PER_VOXEL,
+                max_num_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode],
+            )
 
         points = data_dict['points']
-        voxel_output = voxel_generator.generate(points)
-        if isinstance(voxel_output, dict):
-            voxels, coordinates, num_points = \
-                voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
-        else:
-            voxels, coordinates, num_points = voxel_output
+        voxel_output = self.voxel_generator.generate(points)
+        voxels, coordinates, num_points = voxel_output
 
         if not data_dict['use_lead_xyz']:
             voxels = voxels[..., 3:]  # remove xyz in voxels(N, 3)

pcdet/models/backbones_3d/spconv_backbone.py

 from functools import partial
 
-import spconv
 import torch.nn as nn
 
+from ...utils.spconv_utils import replace_feature, spconv
+
 
 def post_act_block(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=0,
                    conv_type='subm', norm_fn=None):
@@ -50,17 +51,17 @@ class SparseBasicBlock(spconv.SparseModule):
         identity = x
 
         out = self.conv1(x)
-        out.features = self.bn1(out.features)
-        out.features = self.relu(out.features)
+        out = replace_feature(out, self.bn1(out.features))
+        out = replace_feature(out, self.relu(out.features))
 
         out = self.conv2(out)
-        out.features = self.bn2(out.features)
+        out = replace_feature(out, self.bn2(out.features))
 
         if self.downsample is not None:
             identity = self.downsample(x)
 
-        out.features += identity.features
-        out.features = self.relu(out.features)
+        out = replace_feature(out, out.features + identity.features)
+        out = replace_feature(out, self.relu(out.features))
 
         return out
 

pcdet/models/backbones_3d/spconv_unet.py

 from functools import partial
 
-import spconv
 import torch
 import torch.nn as nn
 
+from ...utils.spconv_utils import replace_feature, spconv
 from ...utils import common_utils
 from .spconv_backbone import post_act_block
 
@@ -31,17 +31,17 @@ class SparseBasicBlock(spconv.SparseModule):
         assert x.features.dim() == 2, 'x.features.dim()=%d' % x.features.dim()
 
         out = self.conv1(x)
-        out.features = self.bn1(out.features)
-        out.features = self.relu(out.features)
+        out = replace_feature(out, self.bn1(out.features))
+        out = replace_feature(out, self.relu(out.features))
 
         out = self.conv2(out)
-        out.features = self.bn2(out.features)
+        out = replace_feature(out, self.bn2(out.features))
 
         if self.downsample is not None:
             identity = self.downsample(x)
 
-        out.features += identity
-        out.features = self.relu(out.features)
+        out = replace_feature(out, out.features + identity)
+        out = replace_feature(out, self.relu(out.features))
 
         return out
 
@@ -52,6 +52,7 @@ class UNetV2(nn.Module):
     Reference Paper: https://arxiv.org/abs/1907.03670 (Shaoshuai Shi, et. al)
     From Points to Parts: 3D Object Detection from Point Cloud with Part-aware and Part-aggregation Network
     """
+
     def __init__(self, model_cfg, input_channels, grid_size, voxel_size, point_cloud_range, **kwargs):
         super().__init__()
         self.model_cfg = model_cfg
@@ -134,10 +135,10 @@ class UNetV2(nn.Module):
     def UR_block_forward(self, x_lateral, x_bottom, conv_t, conv_m, conv_inv):
         x_trans = conv_t(x_lateral)
         x = x_trans
-        x.features = torch.cat((x_bottom.features, x_trans.features), dim=1)
+        x = replace_feature(x, torch.cat((x_bottom.features, x_trans.features), dim=1))
         x_m = conv_m(x)
         x = self.channel_reduction(x, x_m.features.shape[1])
-        x.features = x_m.features + x.features
+        x = replace_feature(x, x_m.features + x.features)
         x = conv_inv(x)
         return x
 
@@ -155,7 +156,7 @@ class UNetV2(nn.Module):
         n, in_channels = features.shape
         assert (in_channels % out_channels == 0) and (in_channels >= out_channels)
 
-        x.features = features.view(n, out_channels, -1).sum(dim=2)
+        x = replace_feature(x, features.view(n, out_channels, -1).sum(dim=2))
         return x
 
     def forward(self, batch_dict):

pcdet/models/detectors/detector3d_template.py

@@ -4,6 +4,7 @@ import torch
 import torch.nn as nn
 
 from ...ops.iou3d_nms import iou3d_nms_utils
+from ...utils.spconv_utils import find_all_spconv_keys
 from .. import backbones_2d, backbones_3d, dense_heads, roi_heads
 from ..backbones_2d import map_to_bev
 from ..backbones_3d import pfe, vfe
@@ -325,6 +326,37 @@ class Detector3DTemplate(nn.Module):
             gt_iou = box_preds.new_zeros(box_preds.shape[0])
         return recall_dict
 
+    def _load_state_dict(self, model_state_disk, *, strict=True):
+        state_dict = self.state_dict()  # local cache of state_dict
+
+        spconv_keys = find_all_spconv_keys(self)
+
+        update_model_state = {}
+        for key, val in model_state_disk.items():
+            if key in spconv_keys and key in state_dict and state_dict[key].shape != val.shape:
+                # with different spconv versions, we need to adapt weight shapes for spconv blocks
+                # adapt spconv weights from version 1.x to version 2.x if you used weights from spconv 1.x
+
+                val_native = val.transpose(-1, -2)  # (k1, k2, k3, c_in, c_out) to (k1, k2, k3, c_out, c_in)
+                if val_native.shape == state_dict[key].shape:
+                    val = val_native.contiguous()
+                else:
+                    assert val.shape.__len__() == 5, 'currently only spconv 3D is supported'
+                    val_implicit = val.permute(4, 0, 1, 2, 3)  # (k1, k2, k3, c_in, c_out) to (c_out, k1, k2, k3, c_in)
+                    if val_implicit.shape == state_dict[key].shape:
+                        val = val_implicit.contiguous()
+
+            if key in state_dict and state_dict[key].shape == val.shape:
+                update_model_state[key] = val
+                # logger.info('Update weight %s: %s' % (key, str(val.shape)))
+
+        if strict:
+            self.load_state_dict(update_model_state)
+        else:
+            state_dict.update(update_model_state)
+            self.load_state_dict(state_dict)
+        return state_dict, update_model_state
+
     def load_params_from_file(self, filename, logger, to_cpu=False):
         if not os.path.isfile(filename):
             raise FileNotFoundError
@@ -334,24 +366,17 @@ class Detector3DTemplate(nn.Module):
         checkpoint = torch.load(filename, map_location=loc_type)
         model_state_disk = checkpoint['model_state']
 
-        if 'version' in checkpoint:
-            logger.info('==> Checkpoint trained from version: %s' % checkpoint['version'])
-
-        update_model_state = {}
-        for key, val in model_state_disk.items():
-            if key in self.state_dict() and self.state_dict()[key].shape == model_state_disk[key].shape:
-                update_model_state[key] = val
-                # logger.info('Update weight %s: %s' % (key, str(val.shape)))
+        version = checkpoint.get("version", None)
+        if version is not None:
+            logger.info('==> Checkpoint trained from version: %s' % version)
 
-        state_dict = self.state_dict()
-        state_dict.update(update_model_state)
-        self.load_state_dict(state_dict)
+        state_dict, update_model_state = self._load_state_dict(model_state_disk, strict=False)
 
         for key in state_dict:
             if key not in update_model_state:
                 logger.info('Not updated weight %s: %s' % (key, str(state_dict[key].shape)))
 
-        logger.info('==> Done (loaded %d/%d)' % (len(update_model_state), len(self.state_dict())))
+        logger.info('==> Done (loaded %d/%d)' % (len(update_model_state), len(state_dict)))
 
     def load_params_with_optimizer(self, filename, to_cpu=False, optimizer=None, logger=None):
         if not os.path.isfile(filename):
@@ -363,7 +388,7 @@ class Detector3DTemplate(nn.Module):
         epoch = checkpoint.get('epoch', -1)
         it = checkpoint.get('it', 0.0)
 
-        self.load_state_dict(checkpoint['model_state'])
+        self._load_state_dict(checkpoint['model_state'], strict=True)
 
         if optimizer is not None:
             if 'optimizer_state' in checkpoint and checkpoint['optimizer_state'] is not None:

pcdet/models/roi_heads/partA2_head.py

 import numpy as np
-import spconv
 import torch
 import torch.nn as nn
 
 from ...ops.roiaware_pool3d import roiaware_pool3d_utils
+from ...utils.spconv_utils import spconv
 from .roi_head_template import RoIHeadTemplate
 
 
@@ -192,7 +192,7 @@ class PartA2FCHead(RoIHeadTemplate):
 
         part_features = pooled_part_features[sparse_idx[:, 0], sparse_idx[:, 1], sparse_idx[:, 2], sparse_idx[:, 3]]
         rpn_features = pooled_rpn_features[sparse_idx[:, 0], sparse_idx[:, 1], sparse_idx[:, 2], sparse_idx[:, 3]]
-        coords = sparse_idx.int()
+        coords = sparse_idx.int().contiguous()
         part_features = spconv.SparseConvTensor(part_features, coords, sparse_shape, batch_size_rcnn)
         rpn_features = spconv.SparseConvTensor(rpn_features, coords, sparse_shape, batch_size_rcnn)
 

pcdet/utils/spconv_utils.py

+from typing import Set
+
+try:
+    import spconv.pytorch as spconv
+except:
+    import spconv as spconv
+
+import torch.nn as nn
+
+
+def find_all_spconv_keys(model: nn.Module, prefix="") -> Set[str]:
+    """
+    Finds all spconv keys that need to have weight's transposed
+    """
+    found_keys: Set[str] = set()
+    for name, child in model.named_children():
+        new_prefix = f"{prefix}.{name}" if prefix != "" else name
+
+        if isinstance(child, spconv.conv.SparseConvolution):
+            new_prefix = f"{new_prefix}.weight"
+            found_keys.add(new_prefix)
+
+        found_keys.update(find_all_spconv_keys(child, prefix=new_prefix))
+
+    return found_keys
+
+
+def replace_feature(out, new_features):
+    if "replace_feature" in out.__dir__():
+        # spconv 2.x behaviour
+        return out.replace_feature(new_features)
+    else:
+        out.features = new_features
+        return out

setup.py

@@ -40,7 +40,7 @@ class PostInstallation(install):
 
 
 if __name__ == '__main__':
-    version = '0.3.0+%s' % get_git_commit_number()
+    version = '0.4.0+%s' % get_git_commit_number()
     write_version_to_file(version, 'pcdet/version.py')
 
     setup(
@@ -50,7 +50,7 @@ if __name__ == '__main__':
         install_requires=[
             'numpy',
             'torch>=1.1',
-            'spconv',
+            # 'spconv',  # spconv has different names depending on the cuda version
             'numba',
             'tensorboardX',
             'easydict',