[Feature] Support DGCNN

configs/_base_/models/dgcnn.py

@@ -0,0 +1,29 @@
+# model settings
+model = dict(
+ type='EncoderDecoder3D',
+ backbone=dict(
+ type='DGCNN',
+ in_channels=9, # [xyz, rgb, normal_xyz], modified with dataset
+ num_samples=(20, 20, 20),
+ knn_mods=('D-KNN', 'F-KNN', 'F-KNN'),
+ radius=(None, None, None),
+ gf_channels=((64, 64), (64, 64), (64, )),
+ fa_channels=(1024, ),
+ act_cfg=dict(type='LeakyReLU', negative_slope=0.2),
+ gf_cfg=dict(type='DGCNNGFModule', pool_mod='max')),
+ decode_head=dict(
+ type='DGCNNHead',
+ fp_channels=(1216, 512),
+ channels=256,
+ dropout_ratio=0.5,
+ conv_cfg=dict(type='Conv1d'),
+ norm_cfg=dict(type='BN1d'),
+ act_cfg=dict(type='LeakyReLU', negative_slope=0.2),
+ loss_decode=dict(
+ type='CrossEntropyLoss',
+ use_sigmoid=False,
+ class_weight=None, # modified with dataset
+ loss_weight=1.0)),
+ # model training and testing settings
+ train_cfg=dict(),
+ test_cfg=dict(mode='slide'))

configs/_base_/schedules/seg_cosine_100e.py

@@ -0,0 +1,8 @@
+# optimizer
+# This schedule is mainly used on S3DIS dataset in segmentation task
+optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+lr_config = dict(policy='CosineAnnealing', warmup=None, min_lr=1e-5)
+
+# runtime settings
+runner = dict(type='EpochBasedRunner', max_epochs=100)

configs/dgcnn/README.md

@@ -0,0 +1,43 @@
+# Dynamic Graph CNN for Learning on Point Clouds
+
+## Introduction
+
+<!-- [ALGORITHM] -->
+
+We implement DGCNN and provide the result and checkpoints on S3DIS dataset.
+
+```
+@article{dgcnn,
+ title={Dynamic Graph CNN for Learning on Point Clouds},
+ author={Wang, Yue and Sun, Yongbin and Liu, Ziwei and Sarma, Sanjay E. and Bronstein, Michael M. and Solomon, Justin M.},
+ journal={ACM Transactions on Graphics (TOG)},
+ year={2019}
+}
+```
+
+**Notice**: We follow the implementations in the original DGCNN paper and a PyTorch implementation of DGCNN [code](https://github.com/AnTao97/dgcnn.pytorch).
+
+## Results
+
+### S3DIS
+
+| Method | Split | Lr schd | Mem (GB) | Inf time (fps) | mIoU (Val set) | Download |
+| :-------------------------------------------------------------------------: | :----: | :--------: | :------: | :------------: | :------------: | :----------------------: |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_1 | cosine 100e | 13.1 | | 68.33 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_2 | cosine 100e | 13.1 | | 40.68 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_3 | cosine 100e | 13.1 | | 69.38 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_4 | cosine 100e | 13.1 | | 50.07 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_5 | cosine 100e | 13.1 | | 50.59 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | Area_6 | cosine 100e | 13.1 | | 77.94 | [model](https://download.openmmlab.com/mmdetection3d) &#124; [log](https://download.openmmlab.com/mmdetection3d) |
+| [DGCNN](./dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py) | 6-fold | | | | 59.43 | |
+
+**Notes:**
+
+- We use XYZ+Color+Normalized_XYZ as input in all the experiments on S3DIS dataset.
+- `Area_5` Split means training the model on Area_1, 2, 3, 4, 6 and testing on Area_5.
+- `6-fold` Split means the overall result of 6 different splits (Area_1, Area_2, Area_3, Area_4, Area_5 and Area_6 Splits).
+- Modify `train_area` and `test_area` in the S3DIS dataset's [config](./configs/_base_/datasets/s3dis_seg-3d-13class.py) to set the training areas and testing areas respectively.
+
+## Indeterminism
+
+Since DGCNN testing adopts sliding patch inference which involves random point sampling, and the test script uses fixed random seeds while the random seeds of validation in training are not fixed, the test results may be slightly different from the results reported above.

configs/dgcnn/dgcnn_32x4_consine_100e_s3dis_seg-3d-13class.py

@@ -0,0 +1,24 @@
+_base_ = [
+ '../_base_/datasets/s3dis_seg-3d-13class.py', '../_base_/models/dgcnn.py',
+ '../_base_/schedules/seg_cosine_100e.py', '../_base_/default_runtime.py'
+]
+
+# data settings
+data = dict(samples_per_gpu=32)
+evaluation = dict(interval=2)
+
+# model settings
+model = dict(
+ backbone=dict(in_channels=9), # [xyz, rgb, normalized_xyz]
+ decode_head=dict(
+ num_classes=13, ignore_index=13,
+ loss_decode=dict(class_weight=None)), # S3DIS doesn't use class_weight
+ test_cfg=dict(
+ num_points=4096,
+ block_size=1.0,
+ sample_rate=0.5,
+ use_normalized_coord=True,
+ batch_size=24))
+
+# runtime settings
+checkpoint_config = dict(interval=2)

configs/dgcnn/metafile.yml

@@ -0,0 +1,24 @@
+Collections:
+ - Name: DGCNN
+ Metadata:
+ Training Techniques:
+ - SGD
+ Training Resources: 4x Titan XP GPUs
+ Architecture:
+ - DGCNN
+ Paper: https://arxiv.org/abs/1801.07829
+ README: configs/dgcnn/README.md
+
+Models:
+ - Name: dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py
+ In Collection: DGCNN
+ Config: configs/dgcnn/dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class.py
+ Metadata:
+ Training Data: S3DIS
+ Training Memory (GB): 13.1
+ Results:
+ - Task: 3D Semantic Segmentation
+ Dataset: S3DIS
+ Metrics:
+ mIoU: 59.43
+ Weights: https://download.openmmlab.com/mmdetection3d/v0.1.0_models/dgcnn/dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class/dgcnn_32x4_cosine_100e_s3dis_seg-3d-13class_20210514_143628-4e341a48.pth

mmdet3d/models/backbones/__init__.py

@@ -1,4 +1,5 @@
 from mmdet.models.backbones import SSDVGG, HRNet, ResNet, ResNetV1d, ResNeXt
+from .dgcnn import DGCNN
 from .multi_backbone import MultiBackbone
 from .nostem_regnet import NoStemRegNet
 from .pointnet2_sa_msg import PointNet2SAMSG
@@ -7,5 +8,5 @@
 
 __all__ = [
  'ResNet', 'ResNetV1d', 'ResNeXt', 'SSDVGG', 'HRNet', 'NoStemRegNet',
- 'SECOND', 'PointNet2SASSG', 'PointNet2SAMSG', 'MultiBackbone'
+ 'SECOND', 'DGCNN', 'PointNet2SASSG', 'PointNet2SAMSG', 'MultiBackbone'
 ]

mmdet3d/models/backbones/dgcnn.py

@@ -0,0 +1,93 @@
+from mmcv.runner import BaseModule, auto_fp16
+from torch import nn as nn
+
+from mmdet3d.ops import DGCNNFAModule, build_gf_module
+from mmdet.models import BACKBONES
+
+
+@BACKBONES.register_module()
+class DGCNN(BaseModule):
+ """Backbone network for DGCNN.
+
+ Args:
+ in_channels (int): Input channels of point cloud.
+ num_samples (tuple[int]): The number of samples for knn or ball query
+ in each GF module.
+ knn_mods (tuple[str]): If knn, mod of KNN of each GF module.
+ radius (tuple[float]): Sampling radii of each GF module.
+ gf_channels (tuple[tuple[int]]): Out channels of each mlp in GF module.
+ fa_channels (tuple[int]): Out channels of each mlp in FA module.
+ act_cfg (dict, optional): Config of activation layer.
+ gf_cfg (dict): Config of graph feature module, which may contain the
+ following keys and values:
+
+ - pool_mod (str): Pool method ('max' or 'avg') for GF modules.
+ """
+
+ def __init__(self,
+ in_channels,
+ num_samples=(20, 20, 20),
+ knn_mods=['D-KNN', 'F-KNN', 'F-KNN'],
+ radius=(None, None, None),
+ gf_channels=((64, 64), (64, 64), (64, )),
+ fa_channels=(1024, ),
+ act_cfg=dict(type='ReLU'),
+ gf_cfg=dict(type='DGCNNGFModule', pool_mod='max'),
+ init_cfg=None):
+ super().__init__(init_cfg=init_cfg)
+ self.num_gf = len(gf_channels)
+
+ assert len(num_samples) == len(knn_mods) == len(radius) == len(
+ gf_channels)
+
+ self.GF_modules = nn.ModuleList()
+ gf_in_channel = in_channels * 2
+ skip_channel_list = [gf_in_channel] # input channel list
+
+ for gf_index in range(self.num_gf):
+ cur_gf_mlps = list(gf_channels[gf_index])
+ cur_gf_mlps = [gf_in_channel] + cur_gf_mlps
+ gf_out_channel = cur_gf_mlps[-1]
+
+ self.GF_modules.append(
+ build_gf_module(
+ mlp_channels=cur_gf_mlps,
+ num_sample=num_samples[gf_index],
+ knn_mod=knn_mods[gf_index],
+ radius=radius[gf_index],
+ act_cfg=act_cfg,
+ cfg=gf_cfg))
+ skip_channel_list.append(gf_out_channel)
+ gf_in_channel = gf_out_channel * 2
+
+ fa_in_channel = sum(skip_channel_list[1:])
+ cur_fa_mlps = list(fa_channels)
+ cur_fa_mlps = [fa_in_channel] + cur_fa_mlps
+
+ self.FA_module = DGCNNFAModule(
+ mlp_channels=cur_fa_mlps, act_cfg=act_cfg)
+
+ @auto_fp16(apply_to=('points', ))
+ def forward(self, points):
+ """Forward pass.
+
+ Args:
+ points (torch.Tensor): point coordinates with features,
+ with shape (B, N, in_channels).
+
+ Returns:
+ dict[str, list[torch.Tensor]]: Outputs after GF and FA modules.
+
+ - gf_points (list[torch.Tensor]): Outputs after each GF module.
+ - fa_points (torch.Tensor): Outputs after FA module.
+ """
+ gf_points = [points]
+
+ for i in range(self.num_gf):
+ cur_points = self.GF_modules[i](gf_points[i])
+ gf_points.append(cur_points)
+
+ fa_points = self.FA_module(gf_points)
+
+ ret = dict(gf_points=gf_points, fa_points=fa_points)
+ return ret

mmdet3d/models/decode_heads/__init__.py

@@ -1,4 +1,5 @@
+from .dgcnn_head import DGCNNHead
 from .paconv_head import PAConvHead
 from .pointnet2_head import PointNet2Head
 
-__all__ = ['PointNet2Head', 'PAConvHead']
+__all__ = ['PointNet2Head', 'DGCNNHead', 'PAConvHead']

mmdet3d/models/decode_heads/dgcnn_head.py

@@ -0,0 +1,65 @@
+from mmcv.cnn.bricks import ConvModule
+
+from mmdet3d.ops import DGCNNFPModule
+from mmdet.models import HEADS
+from .decode_head import Base3DDecodeHead
+
+
+@HEADS.register_module()
+class DGCNNHead(Base3DDecodeHead):
+ r"""DGCNN decoder head.
+
+ Decoder head used in `DGCNN <https://arxiv.org/abs/1801.07829>`_.
+ Refer to the
+ `reimplementation code <https://github.com/AnTao97/dgcnn.pytorch>`_.
+
+ Args:
+ fp_channels (tuple[int]): Tuple of mlp channels in FP modules.
+ """
+
+ def __init__(self, fp_channels=(1216, 512), **kwargs):
+ super(DGCNNHead, self).__init__(**kwargs)
+
+ self.FP_module = DGCNNFPModule(
+ mlp_channels=fp_channels, act_cfg=self.act_cfg)
+
+ # https://github.com/charlesq34/pointnet2/blob/master/models/pointnet2_sem_seg.py#L40
+ self.pre_seg_conv = ConvModule(
+ fp_channels[-1],
+ self.channels,
+ kernel_size=1,
+ bias=False,
+ conv_cfg=self.conv_cfg,
+ norm_cfg=self.norm_cfg,
+ act_cfg=self.act_cfg)
+
+ def _extract_input(self, feat_dict):
+ """Extract inputs from features dictionary.
+
+ Args:
+ feat_dict (dict): Feature dict from backbone.
+
+ Returns:
+ torch.Tensor: points for decoder.
+ """
+ fa_points = feat_dict['fa_points']
+
+ return fa_points
+
+ def forward(self, feat_dict):
+ """Forward pass.
+
+ Args:
+ feat_dict (dict): Feature dict from backbone.
+
+ Returns:
+ torch.Tensor: Segmentation map of shape [B, num_classes, N].
+ """
+ fa_points = self._extract_input(feat_dict)
+
+ fp_points = self.FP_module(fa_points)
+ fp_points = fp_points.transpose(1, 2).contiguous()
+ output = self.pre_seg_conv(fp_points)
+ output = self.cls_seg(output)
+
+ return output

mmdet3d/ops/__init__.py

@@ -3,6 +3,8 @@
  sigmoid_focal_loss)
 
 from .ball_query import ball_query
+from .dgcnn_modules import (DGCNNFAModule, DGCNNFPModule, DGCNNGFModule,
+ build_fa_module, build_gf_module)
 from .furthest_point_sample import (Points_Sampler, furthest_point_sample,
  furthest_point_sample_with_dist)
 from .gather_points import gather_points
@@ -33,8 +35,9 @@
  'furthest_point_sample_with_dist', 'three_interpolate', 'three_nn',
  'gather_points', 'grouping_operation', 'group_points', 'GroupAll',
  'QueryAndGroup', 'PointSAModule', 'PointSAModuleMSG', 'PointFPModule',
- 'points_in_boxes_batch', 'get_compiler_version', 'assign_score_withk',
- 'get_compiling_cuda_version', 'Points_Sampler', 'build_sa_module',
+ 'DGCNNFPModule', 'DGCNNGFModule', 'DGCNNFAModule', 'points_in_boxes_batch',
+ 'get_compiler_version', 'assign_score_withk', 'get_compiling_cuda_version',
+ 'Points_Sampler', 'build_sa_module', 'build_gf_module', 'build_fa_module',
  'PAConv', 'PAConvCUDA', 'PAConvSAModuleMSG', 'PAConvSAModule',
  'PAConvCUDASAModule', 'PAConvCUDASAModuleMSG'
 ]

mmdet3d/ops/dgcnn_modules/__init__.py

@@ -0,0 +1,9 @@
+from .builder import build_fa_module, build_gf_module
+from .dgcnn_fa_module import DGCNNFAModule
+from .dgcnn_fp_module import DGCNNFPModule
+from .dgcnn_gf_module import DGCNNGFModule
+
+__all__ = [
+ 'build_fa_module', 'build_gf_module', 'DGCNNFAModule', 'DGCNNFPModule',
+ 'DGCNNGFModule'
+]