pillarnet

pcdet/models/backbones_2d/__init__.py

-from .base_bev_backbone import BaseBEVBackbone
+from .base_bev_backbone import BaseBEVBackbone, BaseBEVBackboneV1
 
 __all__ = {
-    'BaseBEVBackbone': BaseBEVBackbone
+    'BaseBEVBackbone': BaseBEVBackbone,
+    'BaseBEVBackboneV1': BaseBEVBackboneV1
 }

pcdet/models/backbones_2d/base_bev_backbone.py

@@ -110,3 +110,95 @@ class BaseBEVBackbone(nn.Module):
         data_dict['spatial_features_2d'] = x
 
         return data_dict
+
+
+class BaseBEVBackboneV1(nn.Module):
+    def __init__(self, model_cfg, **kwargs):
+        super().__init__()
+        self.model_cfg = model_cfg
+
+        layer_nums = self.model_cfg.LAYER_NUMS
+        num_filters = self.model_cfg.NUM_FILTERS
+        assert len(layer_nums) == len(num_filters) == 2
+
+        num_upsample_filters = self.model_cfg.NUM_UPSAMPLE_FILTERS
+        upsample_strides = self.model_cfg.UPSAMPLE_STRIDES
+        assert len(num_upsample_filters) == len(upsample_strides)
+
+        num_levels = len(layer_nums)
+        self.blocks = nn.ModuleList()
+        self.deblocks = nn.ModuleList()
+        for idx in range(num_levels):
+            cur_layers = [
+                nn.ZeroPad2d(1),
+                nn.Conv2d(
+                    num_filters[idx], num_filters[idx], kernel_size=3,
+                    stride=1, padding=0, bias=False
+                ),
+                nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
+                nn.ReLU()
+            ]
+            for k in range(layer_nums[idx]):
+                cur_layers.extend([
+                    nn.Conv2d(num_filters[idx], num_filters[idx], kernel_size=3, padding=1, bias=False),
+                    nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
+                    nn.ReLU()
+                ])
+            self.blocks.append(nn.Sequential(*cur_layers))
+            if len(upsample_strides) > 0:
+                stride = upsample_strides[idx]
+                if stride >= 1:
+                    self.deblocks.append(nn.Sequential(
+                        nn.ConvTranspose2d(
+                            num_filters[idx], num_upsample_filters[idx],
+                            upsample_strides[idx],
+                            stride=upsample_strides[idx], bias=False
+                        ),
+                        nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3, momentum=0.01),
+                        nn.ReLU()
+                    ))
+                else:
+                    stride = np.round(1 / stride).astype(np.int)
+                    self.deblocks.append(nn.Sequential(
+                        nn.Conv2d(
+                            num_filters[idx], num_upsample_filters[idx],
+                            stride,
+                            stride=stride, bias=False
+                        ),
+                        nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3, momentum=0.01),
+                        nn.ReLU()
+                    ))
+
+        c_in = sum(num_upsample_filters)
+        if len(upsample_strides) > num_levels:
+            self.deblocks.append(nn.Sequential(
+                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1], stride=upsample_strides[-1], bias=False),
+                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
+                nn.ReLU(),
+            ))
+
+        self.num_bev_features = c_in
+
+    def forward(self, data_dict):
+        """
+        Args:
+            data_dict:
+                spatial_features
+        Returns:
+        """
+        spatial_features = data_dict['multi_scale_2d_features']
+
+        x_conv4 = spatial_features['x_conv4']
+        x_conv5 = spatial_features['x_conv5']
+
+        ups = [self.deblocks[0](x_conv4)]
+
+        x = self.blocks[1](x_conv5)
+        ups.append(self.deblocks[1](x))
+
+        x = torch.cat(ups, dim=1)
+        x = self.blocks[0](x)
+
+        data_dict['spatial_features_2d'] = x
+
+        return data_dict

pcdet/models/backbones_3d/__init__.py

 from .pointnet2_backbone import PointNet2Backbone, PointNet2MSG
 from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x
+from .spconv_backbone_2d import PillarBackBone8x, PillarRes18BackBone8x
 from .spconv_backbone_focal import VoxelBackBone8xFocal
 from .spconv_unet import UNetV2
 
@@ -9,5 +10,7 @@ __all__ = {
     'PointNet2Backbone': PointNet2Backbone,
     'PointNet2MSG': PointNet2MSG,
     'VoxelResBackBone8x': VoxelResBackBone8x,
-    'VoxelBackBone8xFocal': VoxelBackBone8xFocal
+    'VoxelBackBone8xFocal': VoxelBackBone8xFocal,
+    'PillarBackBone8x': PillarBackBone8x,
+    'PillarRes18BackBone8x': PillarRes18BackBone8x
 }

pcdet/models/backbones_3d/spconv_backbone_2d.py

+from functools import partial
+
+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):
+
+    if conv_type == 'subm':
+        conv = spconv.SubMConv2d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
+    elif conv_type == 'spconv':
+        conv = spconv.SparseConv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
+                                   bias=False, indice_key=indice_key)
+    elif conv_type == 'inverseconv':
+        conv = spconv.SparseInverseConv2d(in_channels, out_channels, kernel_size, indice_key=indice_key, bias=False)
+    else:
+        raise NotImplementedError
+
+    m = spconv.SparseSequential(
+        conv,
+        norm_fn(out_channels),
+        nn.ReLU(),
+    )
+
+    return m
+
+
+def post_act_block_dense(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, norm_fn=None):
+    m = nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, dilation=dilation, bias=False),
+        norm_fn(out_channels),
+        nn.ReLU(),
+    )
+
+    return m
+
+
+class SparseBasicBlock(spconv.SparseModule):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, norm_fn=None, downsample=None, indice_key=None):
+        super(SparseBasicBlock, self).__init__()
+
+        assert norm_fn is not None
+        bias = norm_fn is not None
+        self.conv1 = spconv.SubMConv2d(
+            inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=bias, indice_key=indice_key
+        )
+        self.bn1 = norm_fn(planes)
+        self.relu = nn.ReLU()
+        self.conv2 = spconv.SubMConv2d(
+            planes, planes, kernel_size=3, stride=stride, padding=1, bias=bias, indice_key=indice_key
+        )
+        self.bn2 = norm_fn(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = replace_feature(out, self.bn1(out.features))
+        out = replace_feature(out, self.relu(out.features))
+
+        out = self.conv2(out)
+        out = replace_feature(out, self.bn2(out.features))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out = replace_feature(out, out.features + identity.features)
+        out = replace_feature(out, self.relu(out.features))
+
+        return out
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, norm_fn=None, downsample=None):
+        super(BasicBlock, self).__init__()
+
+        assert norm_fn is not None
+        bias = norm_fn is not None
+        self.conv1 = nn.Conv2d(inplanes, planes, 3, stride=stride, padding=1, bias=bias)
+        self.bn1 = norm_fn(planes)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(planes, planes, 3, stride=stride, padding=1, bias=bias)
+        self.bn2 = norm_fn(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out = out + identity
+        out = self.relu(out)
+
+        return out
+    
+    
+class PillarBackBone8x(nn.Module):
+    def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
+        super().__init__()
+        self.model_cfg = model_cfg
+        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
+        self.sparse_shape = grid_size[[1, 0]]
+        
+        block = post_act_block
+        dense_block = post_act_block_dense
+
+        self.conv1 = spconv.SparseSequential(
+            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
+            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
+        )
+
+        self.conv2 = spconv.SparseSequential(
+            # [1600, 1408] <- [800, 704]
+            block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
+            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
+            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
+        )
+
+        self.conv3 = spconv.SparseSequential(
+            # [800, 704] <- [400, 352]
+            block(64, 128, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
+            block(128, 128, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
+            block(128, 128, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
+        )
+
+        self.conv4 = spconv.SparseSequential(
+            # [400, 352] <- [200, 176]
+            block(128, 256, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv4', conv_type='spconv'),
+            block(256, 256, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
+            block(256, 256, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
+        )
+        
+        norm_fn = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.01)
+        self.conv5 = nn.Sequential(
+            # [200, 176] <- [100, 88]
+            dense_block(256, 256, 3, norm_fn=norm_fn, stride=2, padding=1),
+            dense_block(256, 256, 3, norm_fn=norm_fn, padding=1),
+            dense_block(256, 256, 3, norm_fn=norm_fn, padding=1),
+        )
+        
+        self.num_point_features = 256
+        self.backbone_channels = {
+            'x_conv1': 32,
+            'x_conv2': 64,
+            'x_conv3': 128,
+            'x_conv4': 256,
+            'x_conv5': 256
+        }
+
+
+    def forward(self, batch_dict):
+        pillar_features, pillar_coords = batch_dict['pillar_features'], batch_dict['pillar_coords']
+        batch_size = batch_dict['batch_size']
+        input_sp_tensor = spconv.SparseConvTensor(
+            features=pillar_features,
+            indices=pillar_coords.int(),
+            spatial_shape=self.sparse_shape,
+            batch_size=batch_size
+        )
+        
+        x_conv1 = self.conv1(input_sp_tensor)
+        x_conv2 = self.conv2(x_conv1)
+        x_conv3 = self.conv3(x_conv2)
+        x_conv4 = self.conv4(x_conv3)
+        x_conv4 = x_conv4.dense()
+        x_conv5 = self.conv5(x_conv4)
+
+        batch_dict.update({
+            'multi_scale_2d_features': {
+                'x_conv1': x_conv1,
+                'x_conv2': x_conv2,
+                'x_conv3': x_conv3,
+                'x_conv4': x_conv4,
+                'x_conv5': x_conv5,
+            }
+        })
+        batch_dict.update({
+            'multi_scale_2d_strides': {
+                'x_conv1': 1,
+                'x_conv2': 2,
+                'x_conv3': 4,
+                'x_conv4': 8,
+                'x_conv5': 16,
+            }
+        })
+
+        return batch_dict
+
+
+class PillarRes18BackBone8x(nn.Module):
+    def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
+        super().__init__()
+        self.model_cfg = model_cfg
+        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
+        self.sparse_shape = grid_size[[1, 0]]
+        
+        block = post_act_block
+        dense_block = post_act_block_dense
+        
+        self.conv1 = spconv.SparseSequential(
+            SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res1'),
+            SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res1'),
+        )
+
+        self.conv2 = spconv.SparseSequential(
+            # [1600, 1408] <- [800, 704]
+            block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
+            SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res2'),
+            SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res2'),
+        )
+
+        self.conv3 = spconv.SparseSequential(
+            # [800, 704] <- [400, 352]
+            block(64, 128, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
+            SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res3'),
+            SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res3'),
+        )
+
+        self.conv4 = spconv.SparseSequential(
+            # [400, 352] <- [200, 176]
+            block(128, 256, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv4', conv_type='spconv'),
+            SparseBasicBlock(256, 256, norm_fn=norm_fn, indice_key='res4'),
+            SparseBasicBlock(256, 256, norm_fn=norm_fn, indice_key='res4'),
+        )
+        
+        norm_fn = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.01)
+        self.conv5 = nn.Sequential(
+            # [200, 176] <- [100, 88]
+            dense_block(256, 256, 3, norm_fn=norm_fn, stride=2, padding=1),
+            BasicBlock(256, 256, norm_fn=norm_fn),
+            BasicBlock(256, 256, norm_fn=norm_fn),
+        )
+
+        self.num_point_features = 256
+        self.backbone_channels = {
+            'x_conv1': 32,
+            'x_conv2': 64,
+            'x_conv3': 128,
+            'x_conv4': 256,
+            'x_conv5': 256
+        }
+
+    def forward(self, batch_dict):
+        pillar_features, pillar_coords = batch_dict['pillar_features'], batch_dict['pillar_coords']
+        batch_size = batch_dict['batch_size']
+        input_sp_tensor = spconv.SparseConvTensor(
+            features=pillar_features,
+            indices=pillar_coords.int(),
+            spatial_shape=self.sparse_shape,
+            batch_size=batch_size
+        )
+        
+        x_conv1 = self.conv1(input_sp_tensor)
+        x_conv2 = self.conv2(x_conv1)
+        x_conv3 = self.conv3(x_conv2)
+        x_conv4 = self.conv4(x_conv3)
+        x_conv4 = x_conv4.dense()
+        x_conv5 = self.conv5(x_conv4)
+
+        # batch_dict.update({
+        #     'encoded_spconv_tensor': out,
+        #     'encoded_spconv_tensor_stride': 8
+        # })
+        batch_dict.update({
+            'multi_scale_2d_features': {
+                'x_conv1': x_conv1,
+                'x_conv2': x_conv2,
+                'x_conv3': x_conv3,
+                'x_conv4': x_conv4,
+                'x_conv5': x_conv5,
+            }
+        })
+        batch_dict.update({
+            'multi_scale_2d_strides': {
+                'x_conv1': 1,
+                'x_conv2': 2,
+                'x_conv3': 4,
+                'x_conv4': 8,
+                'x_conv5': 16,
+            }
+        })
+        
+        return batch_dict

pcdet/models/backbones_3d/vfe/__init__.py

 from .mean_vfe import MeanVFE
 from .pillar_vfe import PillarVFE
 from .dynamic_mean_vfe import DynamicMeanVFE
-from .dynamic_pillar_vfe import DynamicPillarVFE
+from .dynamic_pillar_vfe import DynamicPillarVFE, DynamicPillarPFE
 from .image_vfe import ImageVFE
 from .vfe_template import VFETemplate
 
@@ -12,4 +12,5 @@ __all__ = {
     'ImageVFE': ImageVFE,
     'DynMeanVFE': DynamicMeanVFE,
     'DynPillarVFE': DynamicPillarVFE,
+    'DynamicPillarPFE': DynamicPillarPFE
 }

pcdet/models/backbones_3d/vfe/dynamic_pillar_vfe.py

@@ -140,3 +140,101 @@ class DynamicPillarVFE(VFETemplate):
         batch_dict['pillar_features'] = features
         batch_dict['voxel_coords'] = voxel_coords
         return batch_dict
+
+
+class DynamicPillarPFE(VFETemplate):
+    def __init__(self, model_cfg, num_point_features, voxel_size, grid_size, point_cloud_range, **kwargs):
+        super().__init__(model_cfg=model_cfg)
+
+        self.use_norm = self.model_cfg.USE_NORM
+        self.with_distance = self.model_cfg.WITH_DISTANCE
+        self.use_absolute_xyz = self.model_cfg.USE_ABSLOTE_XYZ
+        self.use_cluster_xyz = self.model_cfg.get('USE_CLUSTER_XYZ', True)
+        if self.use_absolute_xyz:
+            num_point_features += 3
+        if self.use_cluster_xyz:
+            num_point_features += 3
+        if self.with_distance:
+            num_point_features += 1
+
+        self.num_filters = self.model_cfg.NUM_FILTERS
+        assert len(self.num_filters) > 0
+        num_filters = [num_point_features] + list(self.num_filters)
+
+        pfn_layers = []
+        for i in range(len(num_filters) - 1):
+            in_filters = num_filters[i]
+            out_filters = num_filters[i + 1]
+            pfn_layers.append(
+                PFNLayerV2(in_filters, out_filters, self.use_norm, last_layer=(i >= len(num_filters) - 2))
+            )
+        self.pfn_layers = nn.ModuleList(pfn_layers)
+
+        self.voxel_x = voxel_size[0]
+        self.voxel_y = voxel_size[1]
+        self.voxel_z = voxel_size[2]
+        self.x_offset = self.voxel_x / 2 + point_cloud_range[0]
+        self.y_offset = self.voxel_y / 2 + point_cloud_range[1]
+        self.z_offset = self.voxel_z / 2 + point_cloud_range[2]
+
+        self.scale_xy = grid_size[0] * grid_size[1]
+        self.scale_y = grid_size[1]
+
+        self.grid_size = torch.tensor(grid_size[:2]).cuda()
+        self.voxel_size = torch.tensor(voxel_size).cuda()
+        self.point_cloud_range = torch.tensor(point_cloud_range).cuda()
+
+    def get_output_feature_dim(self):
+        return self.num_filters[-1]
+
+    def forward(self, batch_dict, **kwargs):
+        points = batch_dict['points']  # (batch_idx, x, y, z, i, e)
+
+        points_coords = torch.floor(
+            (points[:, [1, 2]] - self.point_cloud_range[[0, 1]]) / self.voxel_size[[0, 1]]).int()
+        mask = ((points_coords >= 0) & (points_coords < self.grid_size[[0, 1]])).all(dim=1)
+        points = points[mask]
+        points_coords = points_coords[mask]
+        points_xyz = points[:, [1, 2, 3]].contiguous()
+
+        merge_coords = points[:, 0].int() * self.scale_xy + \
+                       points_coords[:, 0] * self.scale_y + \
+                       points_coords[:, 1]
+
+        unq_coords, unq_inv, unq_cnt = torch.unique(merge_coords, return_inverse=True, return_counts=True, dim=0)
+
+        f_center = torch.zeros_like(points_xyz)
+        f_center[:, 0] = points_xyz[:, 0] - (points_coords[:, 0].to(points_xyz.dtype) * self.voxel_x + self.x_offset)
+        f_center[:, 1] = points_xyz[:, 1] - (points_coords[:, 1].to(points_xyz.dtype) * self.voxel_y + self.y_offset)
+        f_center[:, 2] = points_xyz[:, 2] - self.z_offset
+
+        features = [f_center]
+        if self.use_absolute_xyz:
+            features.append(points[:, 1:])
+        else:
+            features.append(points[:, 4:])
+
+        if self.use_cluster_xyz:
+            points_mean = torch_scatter.scatter_mean(points_xyz, unq_inv, dim=0)
+            f_cluster = points_xyz - points_mean[unq_inv, :]
+            features.append(f_cluster)
+
+        if self.with_distance:
+            points_dist = torch.norm(points[:, 1:4], 2, dim=1, keepdim=True)
+            features.append(points_dist)
+        features = torch.cat(features, dim=-1)
+
+        for pfn in self.pfn_layers:
+            features = pfn(features, unq_inv)
+
+        # generate voxel coordinates
+        unq_coords = unq_coords.int()
+        pillar_coords = torch.stack((unq_coords // self.scale_xy,
+                                     (unq_coords % self.scale_xy) // self.scale_y,
+                                     unq_coords % self.scale_y,
+                                     ), dim=1)
+        pillar_coords = pillar_coords[:, [0, 2, 1]]
+
+        batch_dict['pillar_features'] = features
+        batch_dict['pillar_coords'] = pillar_coords
+        return batch_dict