single roi aware extractor with unittest

mmdet3d/models/roi_heads/roi_extractors/__init__.py

 from mmdet.models.roi_heads.roi_extractors import SingleRoIExtractor
+from .single_roiaware_extractor import Single3DRoIAwareExtractor
 
-__all__ = ['SingleRoIExtractor']
+__all__ = ['SingleRoIExtractor', 'Single3DRoIAwareExtractor']

mmdet3d/models/roi_heads/roi_extractors/single_roiaware_extractor.py

+import torch
+import torch.nn as nn
+
+from mmdet3d import ops
+from mmdet.models.builder import ROI_EXTRACTORS
+
+
+@ROI_EXTRACTORS.register_module
+class Single3DRoIAwareExtractor(nn.Module):
+    """Point-wise roi-aware Extractor
+
+    Extract Point-wise roi features.
+
+    Args:
+        roi_layer (dict): the config of roi layer
+    """
+
+    def __init__(self, roi_layer=None):
+        super(Single3DRoIAwareExtractor, self).__init__()
+        self.roi_layer = self.build_roi_layers(roi_layer)
+
+    def build_roi_layers(self, layer_cfg):
+        cfg = layer_cfg.copy()
+        layer_type = cfg.pop('type')
+        assert hasattr(ops, layer_type)
+        layer_cls = getattr(ops, layer_type)
+        roi_layers = layer_cls(**cfg)
+        return roi_layers
+
+    def forward(self, feats, coordinate, batch_inds, rois):
+        """Extract point-wise roi features
+
+        Args:
+            feats (FloatTensor): point-wise features with
+                shape (batch, npoints, channels) for pooling
+            coordinate (FloatTensor): coordinate of each point
+            batch_inds (longTensor): indicate the batch of each point
+            rois (FloatTensor): roi boxes with batch indices
+
+        Returns:
+            FloatTensor: pooled features
+        """
+        pooled_roi_feats = []
+        for batch_idx in range(int(batch_inds.max()) + 1):
+            roi_inds = (rois[..., 0].int() == batch_idx)
+            coors_inds = (batch_inds.int() == batch_idx)
+            pooled_roi_feat = self.roi_layer(rois[..., 1:][roi_inds],
+                                             coordinate[coors_inds],
+                                             feats[coors_inds])
+            pooled_roi_feats.append(pooled_roi_feat)
+        pooled_roi_feats = torch.cat(pooled_roi_feats, 0)
+        return pooled_roi_feats

mmdet3d/ops/__init__.py

@@ -2,28 +2,18 @@ from mmdet.ops import (RoIAlign, SigmoidFocalLoss, get_compiler_version,
                        get_compiling_cuda_version, nms, roi_align,
                        sigmoid_focal_loss)
 from .norm import NaiveSyncBatchNorm1d, NaiveSyncBatchNorm2d
+from .roiaware_pool3d import (RoIAwarePool3d, points_in_boxes_cpu,
+                              points_in_boxes_gpu)
 from .sparse_block import (SparseBasicBlock, SparseBasicBlockV0,
                            SparseBottleneck, SparseBottleneckV0)
 from .voxel import DynamicScatter, Voxelization, dynamic_scatter, voxelization
 
 __all__ = [
-    'nms',
-    'soft_nms',
-    'RoIAlign',
-    'roi_align',
-    'get_compiler_version',
-    'get_compiling_cuda_version',
-    'NaiveSyncBatchNorm1d',
-    'NaiveSyncBatchNorm2d',
-    'batched_nms',
-    'Voxelization',
-    'voxelization',
-    'dynamic_scatter',
-    'DynamicScatter',
-    'sigmoid_focal_loss',
-    'SigmoidFocalLoss',
-    'SparseBasicBlockV0',
-    'SparseBottleneckV0',
-    'SparseBasicBlock',
-    'SparseBottleneck',
+    'nms', 'soft_nms', 'RoIAlign', 'roi_align', 'get_compiler_version',
+    'get_compiling_cuda_version', 'NaiveSyncBatchNorm1d',
+    'NaiveSyncBatchNorm2d', 'batched_nms', 'Voxelization', 'voxelization',
+    'dynamic_scatter', 'DynamicScatter', 'sigmoid_focal_loss',
+    'SigmoidFocalLoss', 'SparseBasicBlockV0', 'SparseBottleneckV0',
+    'SparseBasicBlock', 'SparseBottleneck', 'RoIAwarePool3d',
+    'points_in_boxes_gpu', 'points_in_boxes_cpu'
 ]

tests/test_roi_extractors.py

+import pytest
+import torch
+
+from mmdet3d.models.roi_heads.roi_extractors import Single3DRoIAwareExtractor
+
+
+def test_single_roiaware_extractor():
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and torch+cuda')
+
+    roi_layer_cfg = dict(
+        type='RoIAwarePool3d', out_size=4, max_pts_per_voxel=128, mode='max')
+
+    self = Single3DRoIAwareExtractor(roi_layer=roi_layer_cfg)
+    feats = torch.tensor(
+        [[1, 2, 3.3], [1.2, 2.5, 3.0], [0.8, 2.1, 3.5], [1.6, 2.6, 3.6],
+         [0.8, 1.2, 3.9], [-9.2, 21.0, 18.2], [3.8, 7.9, 6.3],
+         [4.7, 3.5, -12.2], [3.8, 7.6, -2], [-10.6, -12.9, -20], [-16, -18, 9],
+         [-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4]],
+        dtype=torch.float32).cuda()
+    coordinate = feats.clone()
+    batch_inds = torch.zeros(feats.shape[0]).cuda()
+    rois = torch.tensor([[0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.3],
+                         [0, -10.0, 23.0, 16.0, 10, 20, 20, 0.5]],
+                        dtype=torch.float32).cuda()
+    # test forward
+    pooled_feats = self(feats, coordinate, batch_inds, rois)
+    assert pooled_feats.shape == torch.Size([2, 4, 4, 4, 3])
+    assert torch.allclose(pooled_feats.sum(),
+                          torch.tensor(51.100).cuda(), 1e-3)