Fix loss name

.isort.cfg

@@ -3,6 +3,6 @@ line_length = 79
 multi_line_output = 0
 known_standard_library = setuptools
 known_first_party = mmdet,mmdet3d
-known_third_party = cv2,load_scannet_data,matplotlib,mmcv,numba,numpy,nuscenes,plyfile,pycocotools,pyquaternion,pytest,scannet_utils,scipy,seaborn,shapely,six,skimage,sunrgbd_utils,terminaltables,torch,torchvision
+known_third_party = cv2,load_scannet_data,matplotlib,mmcv,numba,numpy,nuscenes,plyfile,pycocotools,pyquaternion,pytest,scannet_utils,scipy,seaborn,shapely,skimage,sunrgbd_utils,terminaltables,torch,torchvision
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY

configs/parta2/hv_PartA2_secfpn_2x8_cyclic_80e_kitti-3d-3class.py

+_base_ = ['../_base_/schedules/cyclic_40e.py', '../_base_/default_runtime.py']
+
 # model settings
 voxel_size = [0.05, 0.05, 0.1]
-point_cloud_range = [0, -40, -3, 70.4, 40, 1]  # velodyne coordinates, x, y, z
+point_cloud_range = [0, -40, -3, 70.4, 40, 1]
 
 model = dict(
     type='PartA2',
     voxel_layer=dict(
-        max_num_points=5,  # max_points_per_voxel
+        max_num_points=5,
         point_cloud_range=point_cloud_range,
         voxel_size=voxel_size,
-        max_voxels=(16000, 40000)  # (training, testing) max_coxels
-    ),
+        max_voxels=(16000, 40000)),
     voxel_encoder=dict(type='HardSimpleVFE'),
     middle_encoder=dict(
         type='SparseUNet',
@@ -291,36 +292,8 @@ data = dict(
         modality=input_modality,
         classes=class_names,
         test_mode=True))
-# optimizer
-lr = 0.001  # max learning rate
-optimizer = dict(type='AdamW', lr=lr, betas=(0.95, 0.99), weight_decay=0.01)
-optimizer_config = dict(grad_clip=dict(max_norm=10, norm_type=2))
-lr_config = dict(
-    policy='cyclic',
-    target_ratio=(10, 1e-4),
-    cyclic_times=1,
-    step_ratio_up=0.4)
-momentum_config = dict(
-    policy='cyclic',
-    target_ratio=(0.85 / 0.95, 1),
-    cyclic_times=1,
-    step_ratio_up=0.4)
-checkpoint_config = dict(interval=1)
-evaluation = dict(interval=1)
-# yapf:disable
-log_config = dict(
-    interval=50,
-    hooks=[
-        dict(type='TextLoggerHook'),
-        dict(type='TensorboardLoggerHook')
-    ])
-# yapf:enable
-# runtime settings
-total_epochs = 40
-dist_params = dict(backend='nccl', port=29506)
-log_level = 'INFO'
+
+# Part-A2 uses a different learning rate from what SECOND uses.
+lr = 0.001
+optimizer = dict(lr=lr)
 find_unused_parameters = True
-work_dir = './work_dirs/parta2_secfpn_80e'
-load_from = None
-resume_from = None
-workflow = [('train', 1)]

configs/parta2/hv_PartA2_secfpn_2x8_cyclic_80e_kitti-3d-car.py

-# model settings
+_base_ = './hv_PartA2_secfpn_2x8_cyclic_80e_kitti-3d-3class.py'
+
 voxel_size = [0.05, 0.05, 0.1]
 point_cloud_range = [0, -40, -3, 70.4, 40, 1]  # velodyne coordinates, x, y, z
 
 model = dict(
-    type='PartA2',
-    voxel_layer=dict(
-        max_num_points=5,  # max_points_per_voxel
-        point_cloud_range=point_cloud_range,
-        voxel_size=voxel_size,
-        max_voxels=(16000, 40000)  # (training, testing) max_coxels
-    ),
-    voxel_encoder=dict(type='HardSimpleVFE'),
-    middle_encoder=dict(
-        type='SparseUNet',
-        in_channels=4,
-        sparse_shape=[41, 1600, 1408],
-        order=('conv', 'norm', 'act')),
-    backbone=dict(
-        type='SECOND',
-        in_channels=256,
-        layer_nums=[5, 5],
-        layer_strides=[1, 2],
-        out_channels=[128, 256]),
-    neck=dict(
-        type='SECONDFPN',
-        in_channels=[128, 256],
-        upsample_strides=[1, 2],
-        out_channels=[256, 256]),
     rpn_head=dict(
         type='PartA2RPNHead',
         num_classes=1,
-        in_channels=512,
-        feat_channels=512,
-        use_direction_classifier=True,
         anchor_generator=dict(
+            _delete_=True,
             type='Anchor3DRangeGenerator',
             ranges=[[0, -40.0, -1.78, 70.4, 40.0, -1.78]],
             sizes=[[1.6, 3.9, 1.56]],
             rotations=[0, 1.57],
-            reshape_out=False),
-        diff_rad_by_sin=True,
-        assigner_per_size=True,
-        assign_per_class=True,
-        bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
-        loss_cls=dict(
-            type='FocalLoss',
-            use_sigmoid=True,
-            gamma=2.0,
-            alpha=0.25,
-            loss_weight=1.0),
-        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=2.0),
-        loss_dir=dict(
-            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.2)),
+            reshape_out=False)),
     roi_head=dict(
-        type='PartAggregationROIHead',
         num_classes=1,
-        semantic_head=dict(
-            type='PointwiseSemanticHead',
-            in_channels=16,
-            extra_width=0.2,
-            seg_score_thr=0.3,
-            num_classes=1,
-            loss_seg=dict(
-                type='FocalLoss',
-                use_sigmoid=True,
-                reduction='sum',
-                gamma=2.0,
-                alpha=0.25,
-                loss_weight=1.0),
-            loss_part=dict(
-                type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0)),
-        seg_roi_extractor=dict(
-            type='Single3DRoIAwareExtractor',
-            roi_layer=dict(
-                type='RoIAwarePool3d',
-                out_size=14,
-                max_pts_per_voxel=128,
-                mode='max')),
-        part_roi_extractor=dict(
-            type='Single3DRoIAwareExtractor',
-            roi_layer=dict(
-                type='RoIAwarePool3d',
-                out_size=14,
-                max_pts_per_voxel=128,
-                mode='avg')),
-        bbox_head=dict(
-            type='PartA2BboxHead',
-            num_classes=1,
-            seg_in_channels=16,
-            part_in_channels=4,
-            seg_conv_channels=[64, 64],
-            part_conv_channels=[64, 64],
-            merge_conv_channels=[128, 128],
-            down_conv_channels=[128, 256],
-            bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
-            shared_fc_channels=[256, 512, 512, 512],
-            cls_channels=[256, 256],
-            reg_channels=[256, 256],
-            dropout_ratio=0.1,
-            roi_feat_size=14,
-            with_corner_loss=True,
-            loss_bbox=dict(
-                type='SmoothL1Loss',
-                beta=1.0 / 9.0,
-                reduction='sum',
-                loss_weight=1.0),
-            loss_cls=dict(
-                type='CrossEntropyLoss',
-                use_sigmoid=True,
-                reduction='sum',
-                loss_weight=1.0))))
+        semantic_head=dict(num_classes=1),
+        bbox_head=dict(num_classes=1)))
 # model training and testing settings
 train_cfg = dict(
+    _delete_=True,
     rpn=dict(
-        assigner=dict(  # for Car
+        assigner=dict(
             type='MaxIoUAssigner',
             iou_calculator=dict(type='BboxOverlapsNearest3D'),
             pos_iou_thr=0.6,
@@ -220,71 +129,8 @@ test_pipeline = [
 ]
 
 data = dict(
-    samples_per_gpu=2,
-    workers_per_gpu=2,
-    train=dict(
-        type='RepeatDataset',
-        times=2,
-        dataset=dict(
-            type=dataset_type,
-            data_root=data_root,
-            ann_file=data_root + 'kitti_infos_train.pkl',
-            split='training',
-            pts_prefix='velodyne_reduced',
-            pipeline=train_pipeline,
-            modality=input_modality,
-            classes=class_names,
-            test_mode=False)),
-    val=dict(
-        type=dataset_type,
-        data_root=data_root,
-        ann_file=data_root + 'kitti_infos_val.pkl',
-        split='training',
-        pts_prefix='velodyne_reduced',
-        pipeline=test_pipeline,
-        modality=input_modality,
-        classes=class_names,
-        test_mode=True),
-    test=dict(
-        type=dataset_type,
-        data_root=data_root,
-        ann_file=data_root + 'kitti_infos_val.pkl',
-        split='training',
-        pts_prefix='velodyne_reduced',
-        pipeline=test_pipeline,
-        modality=input_modality,
-        classes=class_names,
-        test_mode=True))
-# optimizer
-lr = 0.001  # max learning rate
-optimizer = dict(type='AdamW', lr=lr, betas=(0.95, 0.99), weight_decay=0.01)
-optimizer_config = dict(grad_clip=dict(max_norm=10, norm_type=2))
-lr_config = dict(
-    policy='cyclic',
-    target_ratio=(10, 1e-4),
-    cyclic_times=1,
-    step_ratio_up=0.4)
-momentum_config = dict(
-    policy='cyclic',
-    target_ratio=(0.85 / 0.95, 1),
-    cyclic_times=1,
-    step_ratio_up=0.4)
-checkpoint_config = dict(interval=1)
-evaluation = dict(interval=1)
-# yapf:disable
-log_config = dict(
-    interval=50,
-    hooks=[
-        dict(type='TextLoggerHook'),
-        dict(type='TensorboardLoggerHook')
-    ])
-# yapf:enable
-# runtime settings
-total_epochs = 40
-dist_params = dict(backend='nccl')
-log_level = 'INFO'
+    train=dict(dataset=dict(pipeline=train_pipeline, classes=class_names)),
+    val=dict(pipeline=test_pipeline, classes=class_names),
+    test=dict(pipeline=test_pipeline, classes=class_names))
+
 find_unused_parameters = True
-work_dir = './work_dirs/parta2_secfpn_80e'
-load_from = None
-resume_from = None
-workflow = [('train', 1)]

mmdet3d/core/__init__.py

 from .anchor import *  # noqa: F401, F403
 from .bbox import *  # noqa: F401, F403
 from .evaluation import *  # noqa: F401, F403
-from .optimizer import *  # noqa: F401, F403
 from .post_processing import *  # noqa: F401, F403
-from .utils import *  # noqa: F401, F403
 from .voxel import *  # noqa: F401, F403

mmdet3d/core/evaluation/__init__.py

-from .class_names import dataset_aliases, get_classes, kitti_classes
 from .indoor_eval import indoor_eval
 from .kitti_utils import kitti_eval, kitti_eval_coco_style
 
-__all__ = [
-    'dataset_aliases', 'get_classes', 'kitti_classes', 'kitti_eval_coco_style',
-    'kitti_eval', 'indoor_eval'
-]
+__all__ = ['kitti_eval_coco_style', 'kitti_eval', 'indoor_eval']

mmdet3d/core/evaluation/class_names.py

-import mmcv
-
-from mmdet.core.evaluation import dataset_aliases
-
-
-def kitti_classes():
-    return [
-        'Car',
-        'Pedestrian',
-        'Cyclist',
-        'Van',
-        'Person_sitting',
-    ]
-
-
-dataset_aliases.update({'kitti': ['KITTI', 'kitti']})
-
-
-def get_classes(dataset):
-    """Get class names of a dataset."""
-    alias2name = {}
-    for name, aliases in dataset_aliases.items():
-        for alias in aliases:
-            alias2name[alias] = name
-
-    if mmcv.is_str(dataset):
-        if dataset in alias2name:
-            labels = eval(alias2name[dataset] + '_classes()')
-        else:
-            raise ValueError('Unrecognized dataset: {}'.format(dataset))
-    else:
-        raise TypeError('dataset must a str, but got {}'.format(type(dataset)))
-    return labels

mmdet3d/core/optimizer/__init__.py

-from .cocktail_constructor import CocktailOptimizerConstructor
-from .cocktail_optimizer import CocktailOptimizer
-
-__all__ = ['CocktailOptimizerConstructor', 'CocktailOptimizer']

mmdet3d/core/optimizer/cocktail_constructor.py

-from mmcv.runner.optimizer import OPTIMIZER_BUILDERS, OPTIMIZERS
-from mmcv.utils import build_from_cfg
-
-from mmdet3d.utils import get_root_logger
-from .cocktail_optimizer import CocktailOptimizer
-
-
-@OPTIMIZER_BUILDERS.register_module()
-class CocktailOptimizerConstructor(object):
-    """Special constructor for cocktail optimizers.
-
-    This constructor constructs cocktail optimizer for multi-modality
-    detectors. It builds separate optimizers for separate branchs for
-    different modalities. More details can be found in the ECCV submission
-    (to be release).
-
-    Attributes:
-        model (:obj:`nn.Module`): The model with parameters to be optimized.
-        optimizer_cfg (dict): The config dict of the optimizer. The keys of
-            the dict are used to search for the corresponding keys in the
-            model, and the value if a dict that really defines the optimizer.
-            See example below for the usage.
-        paramwise_cfg (dict): The dict for paramwise options. This is not
-            supported in the current version. But it should be supported in
-            the future release.
-
-    Example:
-        >>> import torch
-        >>> import torch.nn as nn
-        >>> model = nn.ModuleDict({
-        >>>     'pts': nn.modules.Conv1d(1, 1, 1, bias=False),
-        >>>     'img': nn.modules.Conv1d(1, 1, 1, bias=False)
-        >>> })
-        >>> optimizer_cfg = dict(
-        >>>    pts=dict(type='AdamW', lr=0.001,
-        >>>             weight_decay=0.01, step_interval=1),
-        >>>    img=dict(type='SGD', lr=0.02, momentum=0.9,
-        >>>             weight_decay=0.0001, step_interval=2))
-        >>> optim_builder = CocktailOptimizerConstructor(optimizer_cfg)
-        >>> optimizer = optim_builder(model)
-        >>> print(optimizer)
-        CocktailOptimizer (
-        Update interval: 1
-        AdamW (
-          Parameter Group 0
-              amsgrad: False
-              betas: (0.9, 0.999)
-              eps: 1e-08
-              lr: 0.001
-              weight_decay: 0.01
-          ),
-        Update interval: 2
-        SGD (
-          Parameter Group 0
-              dampening: 0
-              lr: 0.02
-              momentum: 0.9
-              nesterov: False
-              weight_decay: 0.0001
-          ),
-        )
-    """
-
-    def __init__(self, optimizer_cfg, paramwise_cfg=None):
-        if not isinstance(optimizer_cfg, dict):
-            raise TypeError('optimizer_cfg should be a dict',
-                            'but got {}'.format(type(optimizer_cfg)))
-        assert paramwise_cfg is None, \
-            'Parameter wise config is not supported in Cocktail Optimizer'
-        self.optimizer_cfg = optimizer_cfg
-
-    def __call__(self, model):
-        if hasattr(model, 'module'):
-            model = model.module
-
-        optimizer_cfg = self.optimizer_cfg.copy()
-        logger = get_root_logger()
-        keys_prefix = [key_prefix for key_prefix in optimizer_cfg.keys()]
-        keys_params = {key: [] for key in keys_prefix}
-        keys_params_name = {key: [] for key in keys_prefix}
-        keys_optimizer = []
-        for name, param in model.named_parameters():
-            param_group = {'params': [param]}
-            find_flag = False
-            for key in keys_prefix:
-                if key in name:
-                    keys_params[key].append(param_group)
-                    keys_params_name[key].append(name)
-                    find_flag = True
-                    break
-            assert find_flag, 'key {} is not matched to any optimizer'.format(
-                name)
-
-        step_intervals = []
-        for key, single_cfg in optimizer_cfg.items():
-            step_intervals.append(single_cfg.pop('step_interval', 1))
-            single_cfg['params'] = keys_params[key]
-            single_optim = build_from_cfg(single_cfg, OPTIMIZERS)
-            keys_optimizer.append(single_optim)
-            logger.info('{} optimizes key:\n {}\n'.format(
-                single_cfg['type'], keys_params_name[key]))
-
-        cocktail_optimizer = CocktailOptimizer(keys_optimizer, step_intervals)
-        return cocktail_optimizer

mmdet3d/core/optimizer/cocktail_optimizer.py

-from mmcv.runner.optimizer import OPTIMIZERS
-from torch.optim import Optimizer
-
-
-@OPTIMIZERS.register_module()
-class CocktailOptimizer(Optimizer):
-    """Cocktail Optimizer that contains multiple optimizers
-
-    This optimizer applies the cocktail optimzation for multi-modality models.
-
-    Args:
-        optimizers (list[:obj:`torch.optim.Optimizer`]): The list containing
-            different optimizers that optimize different parameters
-        step_intervals (list[int]): Step intervals of each optimizer
-
-    """
-
-    def __init__(self, optimizers, step_intervals=None):
-        self.optimizers = optimizers
-        self.param_groups = []
-        for optimizer in self.optimizers:
-            self.param_groups += optimizer.param_groups
-        if not isinstance(step_intervals, list):
-            step_intervals = [1] * len(self.optimizers)
-        assert len(step_intervals) == len(optimizers), \
-            '"step_intervals" should contain the same number of intervals as' \
-            f'len(optimizers)={len(optimizers)}, got {step_intervals}'
-        self.step_intervals = step_intervals
-        self.num_step_updated = 0
-
-    def __getstate__(self):
-        return {
-            'num_step_updated':
-            self.num_step_updated,
-            'defaults': [optimizer.defaults for optimizer in self.optimizers],
-            'state': [optimizer.state for optimizer in self.optimizers],
-            'param_groups':
-            [optimizer.param_groups for optimizer in self.optimizers],
-        }
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-
-    def __repr__(self):
-        format_string = self.__class__.__name__ + ' (\n'
-        for optimizer, interval in zip(self.optimizers, self.step_intervals):
-            format_string += 'Update interval: {}\n'.format(interval)
-            format_string += optimizer.__repr__().replace('\n', '\n  ') + ',\n'
-        format_string += ')'
-        return format_string
-
-    def state_dict(self):
-        state_dicts = [optimizer.state_dict() for optimizer in self.optimizers]
-        return {
-            'num_step_updated':
-            self.num_step_updated,
-            'state': [state_dict['state'] for state_dict in state_dicts],
-            'param_groups':
-            [state_dict['param_groups'] for state_dict in state_dicts],
-        }
-
-    def load_state_dict(self, state_dict):
-        r"""Loads the optimizer state.
-
-        Arguments:
-            state_dict (dict): optimizer state. Should be an object returned
-                from a call to :meth:`state_dict`.
-        """
-        assert len(state_dict['state']) == len(self.optimizers)
-        assert len(state_dict['param_groups']) == len(self.optimizers)
-        for i, (single_state, single_param_groups) in enumerate(
-                zip(state_dict['state'], state_dict['param_groups'])):
-            single_state_dict = dict(
-                state=single_state, param_groups=single_param_groups)
-            self.optimizers[i].load_state_dict(single_state_dict)
-
-        self.param_groups = []
-        for optimizer in self.optimizers:
-            self.param_groups += optimizer.param_groups
-        self.num_step_updated = state_dict['num_step_updated']
-
-    def zero_grad(self):
-        r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
-        for optimizer in self.optimizers:
-            optimizer.zero_grad()
-
-    def step(self, closure=None):
-        r"""Performs a single optimization step (parameter update).
-
-        Arguments:
-            closure (callable): A closure that reevaluates the model and
-                returns the loss. Optional for most optimizers.
-        """
-        loss = None
-        if closure is not None:
-            loss = closure()
-
-        self.num_step_updated += 1
-        for step_interval, optimizer in zip(self.step_intervals,
-                                            self.optimizers):
-            if self.num_step_updated % step_interval == 0:
-                optimizer.step()
-
-        return loss
-
-    def add_param_group(self, param_group):
-        raise NotImplementedError

mmdet3d/core/utils/__init__.py

-from .dist_utils import DistOptimizerHook, allreduce_grads
-from .misc import tensor2imgs  # merge_batch, merge_hook_batch
-from .misc import multi_apply, unmap
-
-__all__ = [
-    'allreduce_grads',
-    'DistOptimizerHook',
-    'multi_apply',
-    'tensor2imgs',
-    'unmap',  # 'merge_batch', 'merge_hook_batch'
-]

mmdet3d/core/utils/dist_utils.py

-from collections import OrderedDict
-
-import torch.distributed as dist
-from mmcv.runner import OptimizerHook
-from torch._utils import (_flatten_dense_tensors, _take_tensors,
-                          _unflatten_dense_tensors)
-
-
-def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1):
-    if bucket_size_mb > 0:
-        bucket_size_bytes = bucket_size_mb * 1024 * 1024
-        buckets = _take_tensors(tensors, bucket_size_bytes)
-    else:
-        buckets = OrderedDict()
-        for tensor in tensors:
-            tp = tensor.type()
-            if tp not in buckets:
-                buckets[tp] = []
-            buckets[tp].append(tensor)
-        buckets = buckets.values()
-
-    for bucket in buckets:
-        flat_tensors = _flatten_dense_tensors(bucket)
-        dist.all_reduce(flat_tensors)
-        flat_tensors.div_(world_size)
-        for tensor, synced in zip(
-                bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
-            tensor.copy_(synced)
-
-
-def allreduce_grads(params, coalesce=True, bucket_size_mb=-1):
-    grads = [
-        param.grad.data for param in params
-        if param.requires_grad and param.grad is not None
-    ]
-    world_size = dist.get_world_size()
-    if coalesce:
-        _allreduce_coalesced(grads, world_size, bucket_size_mb)
-    else:
-        for tensor in grads:
-            dist.all_reduce(tensor.div_(world_size))
-
-
-class DistOptimizerHook(OptimizerHook):
-
-    def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1):
-        self.grad_clip = grad_clip
-        self.coalesce = coalesce
-        self.bucket_size_mb = bucket_size_mb
-
-    def after_train_iter(self, runner):
-        runner.optimizer.zero_grad()
-        runner.outputs['loss'].backward()
-        # allreduce_grads(runner.model.parameters(), self.coalesce,
-        #                 self.bucket_size_mb)
-        if self.grad_clip is not None:
-            self.clip_grads(runner.model.parameters())
-        runner.optimizer.step()

mmdet3d/core/utils/misc.py

-from functools import partial
-
-import mmcv
-import numpy as np
-import torch
-import torch.nn.functional as F
-from six.moves import map, zip
-
-
-def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
-    num_imgs = tensor.size(0)
-    mean = np.array(mean, dtype=np.float32)
-    std = np.array(std, dtype=np.float32)
-    imgs = []
-    for img_id in range(num_imgs):
-        img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
-        img = mmcv.imdenormalize(
-            img, mean, std, to_bgr=to_rgb).astype(np.uint8)
-        imgs.append(np.ascontiguousarray(img))
-    return imgs
-
-
-def multi_apply(func, *args, **kwargs):
-    pfunc = partial(func, **kwargs) if kwargs else func
-    map_results = map(pfunc, *args)
-    return tuple(map(list, zip(*map_results)))
-
-
-def unmap(data, count, inds, fill=0):
-    """ Unmap a subset of item (data) back to the original set of items (of
-    size count) """
-    if data.dim() == 1:
-        ret = data.new_full((count, ), fill)
-        ret[inds] = data
-    else:
-        new_size = (count, ) + data.size()[1:]
-        ret = data.new_full(new_size, fill)
-        ret[inds, :] = data
-    return ret
-
-
-def merge_batch(data):
-    for key, elems in data.items():
-        if key in ['voxels', 'num_points', 'voxel_labels', 'voxel_centers']:
-            data[key]._data[0] = torch.cat(elems._data[0], dim=0)
-        elif key == 'coors':
-            coors = []
-            for i, coor in enumerate(elems._data[0]):
-                coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
-                coors.append(coor_pad)
-            data[key]._data[0] = torch.cat(coors, dim=0)
-    return data
-
-
-def merge_hook_batch(data):
-    for key, elems in data.items():
-        if key in ['voxels', 'num_points', 'voxel_labels', 'voxel_centers']:
-            data[key] = torch.cat(elems, dim=0)
-        elif key == 'coors':
-            coors = []
-            for i, coor in enumerate(elems):
-                coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
-                coors.append(coor_pad)
-            data[key] = torch.cat(coors, dim=0)
-    return data

mmdet3d/models/dense_heads/anchor3d_head.py

@@ -5,8 +5,8 @@ from mmcv.cnn import bias_init_with_prob, normal_init
 
 from mmdet3d.core import (PseudoSampler, box3d_multiclass_nms, box_torch_ops,
                           boxes3d_to_bev_torch_lidar, build_anchor_generator,
-                          build_assigner, build_bbox_coder, build_sampler,
-                          multi_apply)
+                          build_assigner, build_bbox_coder, build_sampler)
+from mmdet.core import multi_apply
 from mmdet.models import HEADS
 from ..builder import build_loss
 from .train_mixins import AnchorTrainMixin
@@ -326,9 +326,7 @@ class Anchor3DHead(nn.Module, AnchorTrainMixin):
             dir_weights_list,
             num_total_samples=num_total_samples)
         return dict(
-            loss_rpn_cls=losses_cls,
-            loss_rpn_bbox=losses_bbox,
-            loss_rpn_dir=losses_dir)
+            loss_cls=losses_cls, loss_bbox=losses_bbox, loss_dir=losses_dir)
 
     def get_bboxes(self,
                    cls_scores,

mmdet3d/models/dense_heads/parta2_rpn_head.py

@@ -81,6 +81,23 @@ class PartA2RPNHead(Anchor3DHead):
                          diff_rad_by_sin, dir_offset, dir_limit_offset,
                          bbox_coder, loss_cls, loss_bbox, loss_dir)
 
+    def loss(self,
+             cls_scores,
+             bbox_preds,
+             dir_cls_preds,
+             gt_bboxes,
+             gt_labels,
+             input_metas,
+             gt_bboxes_ignore=None):
+        loss_dict = super().loss(cls_scores, bbox_preds, dir_cls_preds,
+                                 gt_bboxes, gt_labels, input_metas,
+                                 gt_bboxes_ignore)
+        # change the loss key names to avoid conflict
+        return dict(
+            loss_rpn_cls=loss_dict['loss_cls'],
+            loss_rpn_bbox=loss_dict['loss_bbox'],
+            loss_rpn_dir=loss_dict['loss_dir'])
+
     def get_bboxes_single(self,
                           cls_scores,
                           bbox_preds,

mmdet3d/models/dense_heads/train_mixins.py

 import numpy as np
 import torch
 
-from mmdet3d.core import box_torch_ops, multi_apply
-from mmdet.core import images_to_levels
+from mmdet3d.core import box_torch_ops
+from mmdet.core import images_to_levels, multi_apply
 
 
 class AnchorTrainMixin(object):

mmdet3d/models/dense_heads/vote_head.py

@@ -4,12 +4,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 from mmcv.cnn import ConvModule
 
-from mmdet3d.core import build_bbox_coder, multi_apply
+from mmdet3d.core import build_bbox_coder
 from mmdet3d.core.post_processing import aligned_3d_nms
 from mmdet3d.models.builder import build_loss
 from mmdet3d.models.losses import chamfer_distance
 from mmdet3d.models.model_utils import VoteModule
 from mmdet3d.ops import PointSAModule, furthest_point_sample
+from mmdet.core import multi_apply
 from mmdet.models import HEADS
 
 

mmdet3d/models/detectors/mvx_two_stage.py

@@ -2,8 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from mmdet3d.core import bbox3d2result, merge_aug_bboxes_3d, multi_apply
+from mmdet3d.core import bbox3d2result, merge_aug_bboxes_3d
 from mmdet3d.ops import Voxelization
+from mmdet.core import multi_apply
 from mmdet.models import DETECTORS
 from .. import builder
 from .base import Base3DDetector

mmdet3d/models/roi_heads/bbox_heads/parta2_bbox_head.py

@@ -4,12 +4,13 @@ import torch.nn as nn
 from mmcv.cnn import ConvModule, normal_init, xavier_init
 
 import mmdet3d.ops.spconv as spconv
-from mmdet3d.core import build_bbox_coder, multi_apply
+from mmdet3d.core import build_bbox_coder
 from mmdet3d.core.bbox import box_torch_ops
 from mmdet3d.models.builder import build_loss
 from mmdet3d.ops import make_sparse_convmodule
 from mmdet3d.ops.iou3d.iou3d_utils import (boxes3d_to_bev_torch_lidar, nms_gpu,
                                            nms_normal_gpu)
+from mmdet.core import multi_apply
 from mmdet.models import HEADS
 
 

mmdet3d/models/roi_heads/mask_heads/pointwise_semantic_head.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from mmdet3d.core import multi_apply
 from mmdet3d.core.bbox import box_torch_ops
 from mmdet3d.models.builder import build_loss
+from mmdet.core import multi_apply
 from mmdet.models import HEADS
 
 

tests/test_heads.py

@@ -105,18 +105,18 @@ def test_anchor3d_head_loss():
 
     losses = self.loss(cls_score, bbox_pred, dir_cls_preds, gt_bboxes,
                        gt_labels, input_metas)
-    assert losses['loss_rpn_cls'][0] > 0
-    assert losses['loss_rpn_bbox'][0] > 0
-    assert losses['loss_rpn_dir'][0] > 0
+    assert losses['loss_cls'][0] > 0
+    assert losses['loss_bbox'][0] > 0
+    assert losses['loss_dir'][0] > 0
 
     # test empty ground truth case
     gt_bboxes = list(torch.empty((2, 0, 7)).cuda())
     gt_labels = list(torch.empty((2, 0)).cuda())
     empty_gt_losses = self.loss(cls_score, bbox_pred, dir_cls_preds, gt_bboxes,
                                 gt_labels, input_metas)
-    assert empty_gt_losses['loss_rpn_cls'][0] > 0
-    assert empty_gt_losses['loss_rpn_bbox'][0] == 0
-    assert empty_gt_losses['loss_rpn_dir'][0] == 0
+    assert empty_gt_losses['loss_cls'][0] > 0
+    assert empty_gt_losses['loss_bbox'][0] == 0
+    assert empty_gt_losses['loss_dir'][0] == 0
 
 
 def test_anchor3d_head_getboxes():