带你玩转 3D 检测和分割（一）：MMDetection3D 整体框架介绍

# MMDetection3D 代码目录结构，展示主要部分
mmdetection3d
   |
   |- configs                    # 配置文件
   |- data                       # 原始数据及预处理后数据文件
   |- mmdet3d 
   |     |- ops                  # cuda 算子（即将迁移到 mmcv 中）
   |     |- core                 # 核心组件
   |     |- datasets             # 数据集相关代码
   |     |- models               # 模型相关代码
   |     |- utils                # 辅助工具
   |     |- ...
   |- tools
   |     |- analysis_tools       # 分析工具，包括可视化、计算flops等
   |     |- data_converter       # 各个数据集预处理转换脚本
   |     |- create_data.py       # 数据预处理入口
   |     |- train.py             # 训练脚本
   |     |- test.py              # 测试脚本
   |     |- ...                      
   |- ...

#=================== mmcv/runner/epoch_based_runner.py ==================
if train_mode:
    outputs = self.model.train_step(data_batch,...)
else:
    outputs = self.model.val_step(data_batch,...)

# 非分布式训练
#=================== mmcv/parallel/data_parallel.py/MMDataParallel ==================
def train_step(self, *inputs, **kwargs):
    if not self.device_ids:
        inputs, kwargs = self.scatter(inputs, kwargs, [-1])
        # 此时才是调用 model 本身的 train_step
        return self.module.train_step(*inputs, **kwargs)
    # 单 gpu 模式
    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
    # 此时才是调用 model 本身的 train_step
    return self.module.train_step(*inputs[0], **kwargs[0])

# val_step 也是的一样逻辑
def val_step(self, *inputs, **kwargs):
    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
    # 此时才是调用 model 本身的 val_step
    return self.module.val_step(*inputs[0], **kwargs[0])

#=================== mmdet/models/detectors/base.py/BaseDetector =============
def train_step(self, data, optimizer):
    # 调用本类自身的 forward 方法
    losses = self(**data)
    # 解析 loss
    loss, log_vars = self._parse_losses(losses)
    # 返回字典对象
    outputs = dict(
        loss=loss, log_vars=log_vars, num_samples=len(data['img_metas']))
    return outputs
    
#=================== mmdet/models/detectors/base.py/Base3DDetector ===========
# Base3DDetector 主要是重写了 forward，改变了模型输入数据的类型，可同时传入点云数据和图片数据，从而满足多模态检测的需求
@auto_fp16(apply_to=('img', 'points'))
def forward(self, return_loss=True, **kwargs):
    if return_loss:
        # 训练模式
        return self.forward_train(**kwargs)
    else:
        # 测试模式
        return self.forward_test(**kwargs)

#============= mmdet/models/detectors/voxelnet.py/VoxelNet ============
def forward_train(self,
                  points,
                  img_metas,
                  gt_bboxes_3d,
                  gt_labels_3d,
                  gt_bboxes_ignore=None):
    # 先进行点云的特征提取  
    x = self.extract_feat(points, img_metas)
    # 主要是调用 bbox_head 内部的 forward_train 方法，得到 head 输出
    outs = self.bbox_head(x)
    loss_inputs = outs + (gt_bboxes_3d, gt_labels_3d, img_metas)
    # 将 head 部分的输出和数据的 label 送入计算 loss
    losses = self.bbox_head.loss(
        *loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
    return losses

#=================== mmdet/models/detectors/base.py/BaseDetector ==================
def _parse_losses(self, losses):
    
    # 返回来的 losses 是一个dict, 我们需要对 loss 进行求和
    log_vars = OrderedDict()
    for loss_name, loss_value in losses.items():
        if isinstance(loss_value, torch.Tensor):
            log_vars[loss_name] = loss_value.mean()
        elif isinstance(loss_value, list):
            log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
        else:
            raise TypeError(
                f'{loss_name} is not a tensor or list of tensors')

    loss = sum(_value for _key, _value in log_vars.items()
               if 'loss' in _key)

    log_vars['loss'] = loss
    for loss_name, loss_value in log_vars.items():
        # reduce loss when distributed training
        if dist.is_available() and dist.is_initialized():
            loss_value = loss_value.data.clone()
            dist.all_reduce(loss_value.div_(dist.get_world_size()))
        log_vars[loss_name] = loss_value.item()

    return loss, log_vars

#=================== mmdet/models/detectors/base.py/Base3DDetector ===========
def forward_test(self, points, img_metas, img=None, **kwargs):
    num_augs = len(points)
    if num_augs != len(img_metas):
        raise ValueError(
            'num of augmentations ({}) != num of image meta ({})'.format(
                len(points), len(img_metas)))
    # 根据 points list 长度判断是 simple_test 还是 aug_test
    if num_augs == 1:
        img = [img] if img is None else img
        return self.simple_test(points[0], img_metas[0], img[0], **kwargs)
    else:
        return self.aug_test(points, img_metas, img, **kwargs)

#============= mmdet/models/detectors/voxelnet.py/VoxelNet ============
def simple_test(self, points, img_metas, imgs=None, rescale=False):
    # 无数据增强测试
    # 提取特征
    x = self.extract_feat(points, img_metas)
    # 调用 head 
    outs = self.bbox_head(x)
    # 根据 head 输出结果生成 bboxes
    bbox_list = self.bbox_head.get_bboxes(
        *outs, img_metas, rescale=rescale)
    # 对检测结果进行格式调整
    bbox_results = [
        bbox3d2result(bboxes, scores, labels)
        for bboxes, scores, labels in bbox_list
    ]
    return bbox_results

def aug_test(self, points, img_metas, imgs=None, rescale=False):
    # 数据增强测试
    feats = self.extract_feats(points, img_metas)

    # 目前只支持单个 sample 的 aug_test
    aug_bboxes = []
    for x, img_meta in zip(feats, img_metas):
        outs = self.bbox_head(x)
        bbox_list = self.bbox_head.get_bboxes(
            *outs, img_meta, rescale=rescale)
        bbox_list = [
            dict(boxes_3d=bboxes, scores_3d=scores, labels_3d=labels)
            for bboxes, scores, labels in bbox_list
        ]
        aug_bboxes.append(bbox_list[0])

    # 将增强后的 bboxes 进行 merge 合并操作
    merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, img_metas,
                                        self.bbox_head.test_cfg)

    return [merged_bboxes]