[Feaute]support spconv2.0 (#1421)

* support spconv2.0

* resolve comments

* try to fix ut

* fix ut

* fix docs

* add spconv2 check in collect.env

* resolve comments

* fix docs

* fix comments

docs/en/getting_started.md

@@ -11,28 +11,28 @@
 The required versions of MMCV, MMDetection and MMSegmentation for different versions of MMDetection3D are as below. Please install the correct version of MMCV, MMDetection and MMSegmentation to avoid installation issues.
 
 | MMDetection3D version |   MMDetection version   | MMSegmentation version |        MMCV version        |
-|:-------------------:|:-------------------:|:-------------------:|:-------------------:|
-| master              | mmdet>=2.19.0, <=3.0.0| mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.4.8, <=1.5.0|
-| v1.0.0rc1           | mmdet>=2.19.0, <=3.0.0| mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.4.8, <=1.5.0|
-| v1.0.0rc0           | mmdet>=2.19.0, <=3.0.0| mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0|
-| 0.18.1              | mmdet>=2.19.0, <=3.0.0| mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0|
-| 0.18.0              | mmdet>=2.19.0, <=3.0.0| mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0|
-| 0.17.3              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.17.2              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.17.1              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.17.0              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.16.0              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.15.0              | mmdet>=2.14.0, <=3.0.0| mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0|
-| 0.14.0              | mmdet>=2.10.0, <=2.11.0| mmseg==0.14.0 | mmcv-full>=1.3.1, <=1.4.0|
-| 0.13.0              | mmdet>=2.10.0, <=2.11.0| Not required  | mmcv-full>=1.2.4, <=1.4.0|
-| 0.12.0              | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.2.4, <=1.4.0|
-| 0.11.0              | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.2.4, <=1.3.0|
-| 0.10.0              | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.2.4, <=1.3.0|
-| 0.9.0               | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.2.4, <=1.3.0|
-| 0.8.0               | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.1.5, <=1.3.0|
-| 0.7.0               | mmdet>=2.5.0, <=2.11.0 | Not required  | mmcv-full>=1.1.5, <=1.3.0|
-| 0.6.0               | mmdet>=2.4.0, <=2.11.0 | Not required  | mmcv-full>=1.1.3, <=1.2.0|
-| 0.5.0               | 2.3.0                  | Not required  | mmcv-full==1.0.5|
+| :-------------------: | :---------------------: | :--------------------: | :------------------------: |
+|        master         | mmdet>=2.19.0, <=3.0.0  | mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.4.8, <=1.5.0  |
+|       v1.0.0rc1       | mmdet>=2.19.0, <=3.0.0  | mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.4.8, <=1.5.0  |
+|       v1.0.0rc0       | mmdet>=2.19.0, <=3.0.0  | mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0 |
+|        0.18.1         | mmdet>=2.19.0, <=3.0.0  | mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0 |
+|        0.18.0         | mmdet>=2.19.0, <=3.0.0  | mmseg>=0.20.0, <=1.0.0 | mmcv-full>=1.3.17, <=1.5.0 |
+|        0.17.3         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.17.2         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.17.1         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.17.0         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.16.0         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.15.0         | mmdet>=2.14.0, <=3.0.0  | mmseg>=0.14.1, <=1.0.0 | mmcv-full>=1.3.8, <=1.4.0  |
+|        0.14.0         | mmdet>=2.10.0, <=2.11.0 |     mmseg==0.14.0      | mmcv-full>=1.3.1, <=1.4.0  |
+|        0.13.0         | mmdet>=2.10.0, <=2.11.0 |      Not required      | mmcv-full>=1.2.4, <=1.4.0  |
+|        0.12.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.2.4, <=1.4.0  |
+|        0.11.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.2.4, <=1.3.0  |
+|        0.10.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.2.4, <=1.3.0  |
+|         0.9.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.2.4, <=1.3.0  |
+|         0.8.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.1.5, <=1.3.0  |
+|         0.7.0         | mmdet>=2.5.0, <=2.11.0  |      Not required      | mmcv-full>=1.1.5, <=1.3.0  |
+|         0.6.0         | mmdet>=2.4.0, <=2.11.0  |      Not required      | mmcv-full>=1.1.3, <=1.2.0  |
+|         0.5.0         |          2.3.0          |      Not required      |      mmcv-full==1.0.5      |
 
 # Installation
 
@@ -192,6 +192,19 @@ you can install it before installing MMCV.
 
 4. Some dependencies are optional. Simply running `pip install -v -e .` will only install the minimum runtime requirements. To use optional dependencies like `albumentations` and `imagecorruptions` either install them manually with `pip install -r requirements/optional.txt` or specify desired extras when calling `pip` (e.g. `pip install -v -e .[optional]`). Valid keys for the extras field are: `all`, `tests`, `build`, and `optional`.
 
+   We have supported spconv2.0. If the user has installed spconv2.0, the code will use spconv2.0 first, which will take up less GPU memory than using the default mmcv spconv. Users can use the following commands to install spconv2.0:
+
+   ```bash
+   pip install cumm-cuxxx
+   pip install spconv-cuxxx
+   ```
+
+   Where xxx is the CUDA version in the environment.
+
+   For example, using CUDA 10.2, the command will be `pip install cumm-cu102 && pip install spconv-cu102`.
+
+   Supported CUDA versions include 10.2, 11.1, 11.3, and 11.4. Users can also install it by building from the source. For more details please refer to [spconv v2.x](https://github.com/traveller59/spconv).
+
    We also support Minkowski Engine as a sparse convolution backend. If necessary please follow original [installation guide](https://github.com/NVIDIA/MinkowskiEngine#installation) or use `pip`:
 
     ```shell

mmdet3d/models/middle_encoders/sparse_encoder.py

 # Copyright (c) OpenMMLab. All rights reserved.
-from mmcv.ops import SparseConvTensor, SparseSequential
+
 from mmcv.runner import auto_fp16
 from torch import nn as nn
 
 from mmdet3d.ops import SparseBasicBlock, make_sparse_convmodule
+from mmdet3d.ops.spconv import IS_SPCONV2_AVAILABLE
 from ..builder import MIDDLE_ENCODERS
 
+if IS_SPCONV2_AVAILABLE:
+    from spconv.pytorch import SparseConvTensor, SparseSequential
+else:
+    from mmcv.ops import SparseConvTensor, SparseSequential
+
 
 @MIDDLE_ENCODERS.register_module()
 class SparseEncoder(nn.Module):

mmdet3d/models/middle_encoders/sparse_unet.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
-from mmcv.ops import SparseConvTensor, SparseSequential
+
+from mmdet3d.ops.spconv import IS_SPCONV2_AVAILABLE
+
+if IS_SPCONV2_AVAILABLE:
+    from spconv.pytorch import SparseConvTensor, SparseSequential
+else:
+    from mmcv.ops import SparseConvTensor, SparseSequential
+
 from mmcv.runner import BaseModule, auto_fp16
 
 from mmdet3d.ops import SparseBasicBlock, make_sparse_convmodule

mmdet3d/ops/sparse_block.py

 # Copyright (c) OpenMMLab. All rights reserved.
 from mmcv.cnn import build_conv_layer, build_norm_layer
-from mmcv.ops import SparseModule, SparseSequential
 from torch import nn
 
 from mmdet.models.backbones.resnet import BasicBlock, Bottleneck
+from .spconv import IS_SPCONV2_AVAILABLE
+
+if IS_SPCONV2_AVAILABLE:
+    from spconv.pytorch import SparseModule, SparseSequential
+else:
+    from mmcv.ops import SparseModule, SparseSequential
+
+
+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
 
 
 class SparseBottleneck(Bottleneck, SparseModule):
@@ -46,21 +60,21 @@ class SparseBottleneck(Bottleneck, SparseModule):
         identity = x.features
 
         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.features = self.relu(out.features)
+        out = replace_feature(out, self.bn2(out.features))
+        out = replace_feature(out, self.relu(out.features))
 
         out = self.conv3(out)
-        out.features = self.bn3(out.features)
+        out = replace_feature(out, self.bn3(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
 
@@ -104,19 +118,18 @@ class SparseBasicBlock(BasicBlock, SparseModule):
         identity = x.features
 
         assert x.features.dim() == 2, f'x.features.dim()={x.features.dim()}'
-
         out = self.conv1(x)
-        out.features = self.norm1(out.features)
-        out.features = self.relu(out.features)
+        out = replace_feature(out, self.norm1(out.features))
+        out = replace_feature(out, self.relu(out.features))
 
         out = self.conv2(out)
-        out.features = self.norm2(out.features)
+        out = replace_feature(out, self.norm2(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
 

mmdet3d/ops/spconv/__init__.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from .overwrite_spconv.write_spconv2 import register_spconv2
+
+try:
+    import spconv
+except ImportError:
+    IS_SPCONV2_AVAILABLE = False
+else:
+    if hasattr(spconv, '__version__') and spconv.__version__ >= '2.0.0':
+        IS_SPCONV2_AVAILABLE = register_spconv2()
+    else:
+        IS_SPCONV2_AVAILABLE = False
+
+__all__ = ['IS_SPCONV2_AVAILABLE']

mmdet3d/ops/spconv/overwrite_spconv/write_spconv2.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import itertools
+
+from mmcv.cnn.bricks.registry import CONV_LAYERS
+from torch.nn.parameter import Parameter
+
+
+def register_spconv2():
+    """This func registers spconv2.0 spconv ops to overwrite the default mmcv
+    spconv ops."""
+    try:
+        from spconv.pytorch import (SparseConv2d, SparseConv3d, SparseConv4d,
+                                    SparseConvTranspose2d,
+                                    SparseConvTranspose3d, SparseInverseConv2d,
+                                    SparseInverseConv3d, SparseModule,
+                                    SubMConv2d, SubMConv3d, SubMConv4d)
+    except ImportError:
+        return False
+    else:
+        CONV_LAYERS._register_module(SparseConv2d, 'SparseConv2d', force=True)
+        CONV_LAYERS._register_module(SparseConv3d, 'SparseConv3d', force=True)
+        CONV_LAYERS._register_module(SparseConv4d, 'SparseConv4d', force=True)
+
+        CONV_LAYERS._register_module(
+            SparseConvTranspose2d, 'SparseConvTranspose2d', force=True)
+        CONV_LAYERS._register_module(
+            SparseConvTranspose3d, 'SparseConvTranspose3d', force=True)
+
+        CONV_LAYERS._register_module(
+            SparseInverseConv2d, 'SparseInverseConv2d', force=True)
+        CONV_LAYERS._register_module(
+            SparseInverseConv3d, 'SparseInverseConv3d', force=True)
+
+        CONV_LAYERS._register_module(SubMConv2d, 'SubMConv2d', force=True)
+        CONV_LAYERS._register_module(SubMConv3d, 'SubMConv3d', force=True)
+        CONV_LAYERS._register_module(SubMConv4d, 'SubMConv4d', force=True)
+        SparseModule._load_from_state_dict = _load_from_state_dict
+        SparseModule._save_to_state_dict = _save_to_state_dict
+        return True
+
+
+def _save_to_state_dict(self, destination, prefix, keep_vars):
+    """Rewrite this func to compat the convolutional kernel weights between
+    spconv 1.x in MMCV and 2.x in spconv2.x.
+
+    Kernel weights in MMCV spconv has shape in (D,H,W,in_channel,out_channel) ,
+    while those in spcon2.x is in (out_channel,D,H,W,in_channel).
+    """
+    for name, param in self._parameters.items():
+        if param is not None:
+            param = param if keep_vars else param.detach()
+            if name == 'weight':
+                dims = list(range(1, len(param.shape))) + [0]
+                param = param.permute(*dims)
+            destination[prefix + name] = param
+    for name, buf in self._buffers.items():
+        if buf is not None and name not in self._non_persistent_buffers_set:
+            destination[prefix + name] = buf if keep_vars else buf.detach()
+
+
+def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                          missing_keys, unexpected_keys, error_msgs):
+    """Rewrite this func to compat the convolutional kernel weights between
+    spconv 1.x in MMCV and 2.x in spconv2.x.
+
+    Kernel weights in MMCV spconv has shape in (D,H,W,in_channel,out_channel) ,
+    while those in spcon2.x is in (out_channel,D,H,W,in_channel).
+    """
+    for hook in self._load_state_dict_pre_hooks.values():
+        hook(state_dict, prefix, local_metadata, strict, missing_keys,
+             unexpected_keys, error_msgs)
+
+    local_name_params = itertools.chain(self._parameters.items(),
+                                        self._buffers.items())
+    local_state = {k: v.data for k, v in local_name_params if v is not None}
+
+    for name, param in local_state.items():
+        key = prefix + name
+        if key in state_dict:
+            input_param = state_dict[key]
+
+            # Backward compatibility: loading 1-dim tensor from
+            # 0.3.* to version 0.4+
+            if len(param.shape) == 0 and len(input_param.shape) == 1:
+                input_param = input_param[0]
+            dims = [len(input_param.shape) - 1] + list(
+                range(len(input_param.shape) - 1))
+            input_param = input_param.permute(*dims)
+            if input_param.shape != param.shape:
+                # local shape should match the one in checkpoint
+                error_msgs.append(
+                    f'size mismatch for {key}: copying a param with '
+                    f'shape {key, input_param.shape} from checkpoint,'
+                    f'the shape in current model is {param.shape}.')
+                continue
+
+            if isinstance(input_param, Parameter):
+                # backwards compatibility for serialized parameters
+                input_param = input_param.data
+            try:
+                param.copy_(input_param)
+            except Exception:
+                error_msgs.append(
+                    f'While copying the parameter named "{key}", whose '
+                    f'dimensions in the model are {param.size()} and whose '
+                    f'dimensions in the checkpoint are {input_param.size()}.')
+        elif strict:
+            missing_keys.append(key)
+
+    if strict:
+        for key, input_param in state_dict.items():
+            if key.startswith(prefix):
+                input_name = key[len(prefix):]
+                input_name = input_name.split(
+                    '.', 1)[0]  # get the name of param/buffer/child
+                if input_name not in self._modules \
+                        and input_name not in local_state:
+                    unexpected_keys.append(key)

mmdet3d/utils/collect_env.py

@@ -5,6 +5,7 @@ from mmcv.utils import get_git_hash
 import mmdet
 import mmdet3d
 import mmseg
+from mmdet3d.ops.spconv import IS_SPCONV2_AVAILABLE
 
 
 def collect_env():
@@ -13,7 +14,7 @@ def collect_env():
     env_info['MMDetection'] = mmdet.__version__
     env_info['MMSegmentation'] = mmseg.__version__
     env_info['MMDetection3D'] = mmdet3d.__version__ + '+' + get_git_hash()[:7]
-
+    env_info['spconv2.0'] = IS_SPCONV2_AVAILABLE
     return env_info
 
 

requirements/optional.txt

 open3d
+spconv
 waymo-open-dataset-tf-2-1-0==1.2.0