PointPillars检测模型训练

这篇教程主要是告诉大家如何利用HAT在雷达点云数据集 KITTI-3DObject 上从头开始训练一个 PointPillars 模型,包括浮点、量化和定点模型。

训练流程

如果你只是想简单的把 PointPillars 的模型训练起来,那么可以首先阅读一下这一章的内容。 和其他任务一样,对于所有的训练,评测任务,HAT统一采用 tools + config 的形式来完成。 在准备好原始数据集之后,可以通过下面的流程,方便地完成整个训练的流程。

数据集准备

在开始训练模型之前,第一步是需要准备好数据集,我们在KITTI官网下载 3DObject据集 , 包括4个文件:

  1. left color images of object data set
  2. velodyne point clouds
  3. camera calibration matrices of object data set
  4. taining labels of object data set

下载上述4个文件后,解压并按照如下方式组织文件夹结构:

├── tmp_data
│   ├── kitti3d
│   │   ├── testing
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── velodyne
│   │   ├── training
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── label_2
│   │   │   ├── velodyne

为了创建KITTI点云数据,首先需要加载原始的点云数据并生成相关的包含目标标签和标注框的数据标注文件,同时还需要为KITTI数据集生成每个单独的训练目标的点云数据,并将其存储在 data/kitti/gt_database.bin 格式的文件中,此外,需要为训练数据或者验证数据生成 .pkl 格式的包含数据信息的文件。随后,通过运行下面的命令来创建KITTI数据:

mkdir ./tmp_data/kitti3d/ImageSets
# 从社区下载数据集划分文件
wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/test.txt --no-check-certificate --content-disposition -O ./tmp_data/kitti3d/ImageSets/test.txt
wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/train.txt --no-check-certificate --content-disposition -O ./tmp_data/kitti3d/ImageSets/train.txt
wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/val.txt --no-check-certificate --content-disposition -O ./tmp_data/kitti3d/ImageSets/val.txt
wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/trainval.txt --no-check-certificate --content-disposition -O ./tmp_data/kitti3d/ImageSets/trainval.txt
python3 tools/create_data.py --dataset "kitti3d" --root-dir "./tmp_data/kitti3d"

执行上述命令后,生成的文件目录如下:

├── tmp_data
│   ├──── kitti3d
│   │   ├── ImageSets
│   │   │   ├── test.txt
│   │   │   ├── train.txt
│   │   │   ├── trainval.txt
│   │   │   ├── val.txt
│   │   ├── testing
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── velodyne
│   │   │   ├── velodyne_reduced        # 新生成的 velodyne_reduced
│   │   ├── training
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── label_2
│   │   │   ├── velodyne
│   │   │   ├── velodyne_reduced        # 新生成的 velodyne_reduced
│   │   ├── kitti3d_gt_database           # 新生成的 kitti_gt_database
│   │   │   ├── xxxxx.bin
│   │   ├── kitti3d_infos_train.pkl       # 新生成的 kitti_infos_train.pkl
│   │   ├── kitti3d_infos_val.pkl         # 新生成的 kitti_infos_val.pkl
│   │   ├── kitti3d_dbinfos_train.pkl     # 新生成的 kitti_dbinfos_train.pkl
│   │   ├── kitti3d_infos_test.pkl        # 新生成的 kitti_infos_test.pkl
│   │   ├── kitti3d_infos_trainval.pkl    # 新生成的 kitti_infos_trainval.pkl

同时,为了提升训练的速度,我们对数据信息文件做了一个打包,将其转换成lmdb格式的数据集。只需要运行下面的脚本,就可以成功实现转换:

python3 tools/datasets/kitti3d_packer.py --src-data-dir ./tmp_data/kitti3d/ --target-data-dir ./tmp_data/kitti3d --split-name train --pack-type lmdb
python3 tools/datasets/kitti3d_packer.py --src-data-dir ./tmp_data/kitti3d/ --target-data-dir ./tmp_data/kitti3d --split-name val --pack-type lmdb

上面这两条命令分别对应着转换训练数据集和验证数据集,打包完成之后,data目录下的文件结构应该如下所示:

├── tmp_data
│   ├──── kitti3d
│   │   ├── pack_data       # 新生成的 lmdb
│   │   │   ├── train
│   │   │   ├── val
│   │   ├── ImageSets
│   │   │   ├── test.txt
│   │   │   ├── train.txt
│   │   │   ├── trainval.txt
│   │   │   ├── val.txt
│   │   ├── testing
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── velodyne
│   │   │   ├── velodyne_reduced
│   │   ├── training
│   │   │   ├── calib
│   │   │   ├── image_2
│   │   │   ├── label_2
│   │   │   ├── velodyne
│   │   │   ├── velodyne_reduced
│   │   ├── kitti3d_gt_database
│   │   │   ├── xxxxx.bin
│   │   ├── kitti3d_infos_train.pkl
│   │   ├── kitti3d_infos_val.pkl
│   │   ├── kitti3d_dbinfos_train.pkl
│   │   ├── kitti3d_infos_test.pkl
│   │   ├── kitti3d_infos_trainval.pkl

train_lmdbval_lmdb 就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集,kitti3d_gt_databasekitti3d_dbinfos_train.pkl 是训练是用于采样的样本。

模型训练

数据集准备好之后,就可以开始训练浮点型的PointPillars检测网络了。在网络训练开始之前,你可以使用以下命令先测试一下网络的计算量和参数数量:

python3 tools/calops.py --config configs/detection/pointpillars/pointpillars_kitti_car.py

如果你只是单纯的想启动这样的训练任务,只需要运行下面的命令就可以:

python3 tools/train.py --stage "float" --config configs/detection/pointpillars/pointpillars_kitti_car.py

由于HAT算法包使用了一种巧妙的注册机制,使得每一个训练任务都可以按照这种train.py加上config配置文件的形式启动。 train.py 是统一的训练脚本,与任务无关,我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的config配置文件里面。 config文件里面提供了模型构建、数据读取等关键的dict。

导出定点模型

完成量化训练后,便可以开始导出定点模型。可以通过下面命令来导出:

python3 tools/export_hbir.py --config configs/detection/pointpillars/pointpillars_kitti_car.py

模型验证

模型训练完成之后,我们还可以验证训练出来的模型性能。由于我们提供了float、calibration和qat三阶段的训练过程,相应的我们可以验证这三个阶段训练出来的模型性能,只需要相应的运行以下三条命令即可:

python3 tools/predict.py --stage "float" --config configs/detection/pointpillars/pointpillars_kitti_car.py
python3 tools/predict.py --stage "calibration" --config configs/detection/pointpillars/pointpillars_kitti_car.py
python3 tools/predict.py --stage "qat" --config configs/detection/pointpillars/pointpillars_kitti_car.py

定点模型精度验证可使用下面命令,但需要注意是必须要先导出hbir:

python3 tools/predict.py --stage "int_infer" --config configs/detection/pointpillars/pointpillars_kitti_car.py

这个显示出来的精度才是最终的int8模型的真正精度,当然这个精度和qat验证阶段的精度应该是保持十分接近的。

模型推理和结果可视化

HAT提供了 infer_hbir.py 脚本对各阶段训练好的模型的推理结果进行可视化展示:

python3 tools/infer_hbir.py --config configs/detection/pointpillars/pointpillars_kitti_car.py --model-inputs ${model_inputs} --save-path ${save_path} --use-dataset

python3 tools/infer_hbir.py --config configs/detection/pointpillars/pointpillars_kitti_car.py --model-inputs ${model_inputs} --save-path ${save_path}

定点模型检查和编译

在HAT中集成的量化训练工具链主要是为了地平线的计算平台准备的,因此,对于量化模型的检查和编译是必须的。 我们在HAT中提供了模型检查的接口,可以在定义好量化模型之后,先检查能否在 BPU 上正常运行:

python3 tools/model_checker.py --config configs/detection/pointpillars/pointpillars_kitti_car.py

在训练完成之后,可以使用 compile_perf_hbir 的工具用来将量化模型编译成可以上板运行的 hbm 文件,同时该工具也能预估在 BPU 上的运行性能,可以采用以下脚本:

python3 tools/compile_perf_hbir.py --config configs/detection/pointpillars/pointpillars_kitti_car.py

训练细节

在这个说明中,我们对模型训练需要注意的一些事项进行说明,主要为 config 的一些相关设置。

模型构建

PointPillars 的网络结构可以参考 论文, 这里不做详细介绍。我们通过在config配置文件中定义 model 这样的一个dict型变量,就可以方便的实现对模型的定义和修改。

model = dict(
    type="PointPillarsDetector",
    feature_map_shape=get_feature_map_size(pc_range, voxel_size),
    pre_process=dict(
        type="PointPillarsPreProcess",
        pc_range=pc_range,
        voxel_size=voxel_size,
        max_voxels_num=max_voxels_num,
        max_points_in_voxel=max_points_in_voxel,
    ),
    reader=dict(
        type="PillarFeatureNet",
        num_input_features=4,
        num_filters=(64,),
        with_distance=False,
        pool_size=(1, max_points_in_voxel),
        voxel_size=voxel_size,
        pc_range=pc_range,
        bn_kwargs=norm_cfg,
        quantize=True,
        use_4dim=True,
    ),
    backbone=dict(
        type="PointPillarScatter",
        num_input_features=64,
        use_horizon_pillar_scatter=True,
        quantize=True,
    ),
    neck=dict(
        type="SequentialBottleNeck",
        layer_nums=[3, 5, 5],
        ds_layer_strides=[2, 2, 2],
        ds_num_filters=[64, 128, 256],
        us_layer_strides=[1, 2, 4],
        us_num_filters=[128, 128, 128],
        num_input_features=64,
        bn_kwargs=norm_cfg,
        use_tconv=True,
        use_secnet=True,
        quantize=True,
    ),
    head=dict(
        type="PointPillarHead",
        num_classes=len(class_names),
        in_channels=sum([128, 128, 128]),
        use_direction_classifier=True,
    ),
    anchor_generator=dict(
        type="Anchor3DGeneratorStride",
        anchor_sizes=[[1.6, 3.9, 1.56]],
        anchor_strides=[[0.32, 0.32, 0.0]],
        anchor_offsets=[[0.16, -39.52, -1.78]],
        rotations=[[0, 1.57]],
        class_names=class_names,
        match_thresholds=[0.6],
        unmatch_thresholds=[0.45],
    ),
    targets=dict(
        type="LidarTargetAssigner",
        box_coder=dict(
            type="GroundBox3dCoderTorch",
            n_dim=7,
        ),
        class_names=class_names,
        positive_fraction=-1,
        region_similarity_calculator=dict(type="NearestIouSimilarity"),
    ),
    loss=dict(
        type="PointPillarsLoss",
        num_classes=len(class_names),
        loss_cls=dict(
            type="SigmoidFocalLoss",
            alpha=0.25,
            gamma=2.0,
            loss_weight=1.0,
        ),
        loss_bbox=dict(
            type="WeightedSmoothL1Loss",
            sigma=3.0,
            code_weights=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            codewise=True,
            loss_weight=2.0,
        ),
        loss_dir=dict(
            type="WeightedSoftmaxClassificationLoss",
            name="direction_classifier",
            loss_weight=0.2,
        ),
    ),
    postprocess=dict(
        type="PointPillarsPostProcess",
        num_classes=len(class_names),
        box_coder=dict(
            type="GroundBox3dCoderTorch",
            n_dim=7,
        ),
        use_direction_classifier=True,
        num_direction_bins=2,
        # test_cfg
        use_rotate_nms=False,
        nms_pre_max_size=1000,
        nms_post_max_size=300,
        nms_iou_threshold=0.5,
        score_threshold=0.4,
        post_center_limit_range=[0, -39.68, -5, 69.12, 39.68, 5],
        max_per_img=100,
    ),
)

其中,model 下面的 type 表示定义的模型名称,剩余的变量表示模型的其他组成部分。这样定义模型的好处在于我们可以很方便的替换我们想要的结构。 训练脚本在启动之后,会调用 build_model 接口,将这样一个dict类型的model变成类型为 torch.nn.Module 类型的model。

数据增强

model 的定义一样,数据增强的流程是通过在config配置文件中定义 data_loaderval_data_loader 这两个dict来实现的,分别对应着训练集和验证集的处理流程。以 data_loader 为例:

dataset = dict(
  type="Kitti3D",
  data_path="./tmp_data/kitti3d/train_lmdb",
  transforms=[
      dict(
          type="ObjectSample",
          class_names=class_names,
          remove_points_after_sample=False,
          db_sampler=db_sampler,
      ),
      dict(
          type="ObjectNoise",
          gt_rotation_noise=[-0.15707963267, 0.15707963267],
          gt_loc_noise_std=[0.25, 0.25, 0.25],
          global_random_rot_range=[0, 0],
          num_try=100,
          class_names=class_names,
      ),
      dict(
          type="PointRandomFlip",
          probability=0.5,
      ),
      dict(
          type="PointGlobalRotation",
          rotation=[-0.78539816, 0.78539816],
      ),
      dict(
          type="PointGlobalScaling",
          min_scale=0.95,
          max_scale=1.05,
      ),
      dict(
          type="ShufflePoints",
      ),
      dict(
          type="ObjectRangeFilter",
          point_cloud_range=pc_range,
      ),
      dict(type="Reformat"),
  ],
)

data_loader = dict(
  type=torch.utils.data.DataLoader,
  dataset=dataset,
  sampler=dict(type=torch.utils.data.DistributedSampler),
  batch_size=batch_size_per_gpu,
  shuffle=False,
  num_workers=1,
  pin_memory=True,
  collate_fn=hat.data.collates.collate_kitti3d,
)

其中type直接用的pytorch自带的接口 torch.utils.data.DataLoader,表示的是将 batch_size 大小的图片组合到一起。 这里面唯一需要关注的可能是 dataset 这个变量,data_path 路径也就是我们在第一部分数据集准备中提到的路径。transforms 下面包含着一系列的数据增强。val_data_loader 中只有除了点云Pillar化(Voxelization)和Reformat。 你也可以通过在 transforms 中插入新的dict实现自己希望的数据增强操作。

训练策略

为了训练一个精度高的模型,好的训练策略是必不可少的。对于每一个训练任务而言,相应的训练策略同样都定义在其中的config文件中,从 float_trainer 这个变量就可以看出来。

float_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        betas=(0.95, 0.99),
        lr=5e-4,
        weight_decay=0.1,
    ),
    batch_processor=batch_processor,
    num_epochs=160,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(
            type="CyclicLrUpdater",
            target_ratio=(10, 1e-4),
            cyclic_times=1,
            step_ratio_up=0.4,
            step_log_interval=50,
        ),
        dict(
            type="CyclicOptimParamUpdater",
            param_name="betas",
            target_ratio=(0.85 / 0.95, 1),
            cyclic_times=1,
            step_ratio_up=0.4,
            step_log_interval=50,
        ),
        dict(
            type="GradClip",
            max_norm=35,
            norm_type=2,
        ),
        val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="Kitti3DMetricDet",
        compute_aos=True,
        current_classes=class_names,
        difficultys=[0, 1, 2],
    ),
)

float_trainer 从大局上定义了我们的训练方式,包括使用多卡分布式训练(distributed_data_parallel_trainer),模型训练的epoch次数,以及优化器的选择。 同时 callbacks 中体现了模型在训练过程中使用到的小策略以及您想实现的操作,包括学习率的变换方式(CyclicLrUpdater),在训练过程中验证模型的指标(Validation),以及保存(Checkpoint)模型的操作。当然,如果你有自己希望模型在训练过程中实现的操作,也可以按照这种dict的方式添加。 float_trainer 负责将整个训练的逻辑给串联起来,其中也会负责模型的pretrain。

注解

如果需要复现精度,config中的训练策略最好不要修改。否则可能会有意外的训练情况出现。

通过上面的介绍,你应该对config文件的功能有了一个比较清楚的认识。然后通过前面提到的训练脚本,就可以训练一个高精度的纯浮点的检测模型。 当然训练一个好的检测模型不是我们最终的目的,它只是做为一个pretrain为我们后面训练定点模型服务的。

量化模型训练

当我们有了纯浮点模型之后,就可以开始训练相应的定点模型了。和浮点训练的方式一样,我们只需要通过运行下面的脚本就可以训练定点模型了:

python3 tools/train.py --stage "calibration" --config configs/detection/pointpillars/pointpillars_kitti_car.py

python3 tools/train.py --stage "qat" --config configs/detection/pointpillars/pointpillars_kitti_car.py

可以看到,我们的配置文件没有改变,只改变了 stage 的类型。此时我们使用的训练策略来自于config文件中的 qat_trainercalibration_trainer

calibration_trainer = dict(
    type="Calibrator",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
            ),
            dict(
                type="Float2Calibration",
                convert_mode=convert_mode,
                example_data_loader=calibration_example_data_loader,
                qconfig_setter=calib_qconfig_setter,
            ),
        ],
    ),
    data_loader=calibration_data_loader,
    batch_processor=calibration_batch_processor,
    num_steps=calibration_step,
    device=None,
    callbacks=[
        stat_callback,
        val_callback,
        ckpt_callback,
    ],
    val_metrics=dict(
        type="Kitti3DMetricDet",
        compute_aos=True,
        current_classes=class_names,
        difficultys=[0, 1, 2],
    ),
    log_interval=calibration_step / 10,
)

qat_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        qconfig_params=dict(
            activation_qat_qkwargs=dict(
                averaging_constant=0,
            ),
            weight_qat_qkwargs=dict(
                averaging_constant=1,
            ),
        ),
        converters=[
            dict(
                type="Float2QAT",
                convert_mode=convert_mode,
                example_data_loader=copy.deepcopy(
                    calibration_example_data_loader
                ),
                qconfig_setter=qat_qconfig_setter,
            ),
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "calibration-checkpoint-best.pth.tar"
                ),
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.SGD,
        params={"weight": dict(weight_decay=0.0)},
        lr=2e-4,
        momentum=0.9,
    ),
    batch_processor=batch_processor,
    num_epochs=50,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(
            type="CyclicLrUpdater",
            target_ratio=(10, 1e-4),
            cyclic_times=1,
            step_ratio_up=0.4,
            step_log_interval=50,
        ),
        val_callback,
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="Kitti3DMetricDet",
        compute_aos=True,
        current_classes=class_names,
        difficultys=[0, 1, 2],
    ),
)

quantize参数的值不同

当我们训练量化模型的时候,需要设置模型的qconfig,浮点模型会按照qconfig的设置被转换成对应的量化模型,qconfig设置如下:

def get_custom_qconfig(is_calib):
    my_qconfig_setter = QconfigSetter(
        reference_qconfig=get_qconfig(  # 1. 主要用于获取 observer
            observer=(
                observer_v2.MSEObserver
                if is_calib
                else observer_v2.MinMaxObserver
            )
        ),
        templates=[  # 2. 模板,仅关注 dtype,按次序生效,因此模板的顺序很重要
            # 2.1. 基础配置部分
            ModuleNameTemplate({"": qint8}),
            MatmulDtypeTemplate(
                input_dtypes=[qint8, qint16],
            ),
            ConvDtypeTemplate(
                input_dtype=qint8,
                weight_dtype=qint8,
            ),
        ],
        enable_optimize=True,
        save_dir="./qconfig_setting",
        custom_qconfig_mapping=None,
    )
    return my_qconfig_setter


qat_qconfig_setter = get_custom_qconfig(False)
calib_qconfig_setter = get_custom_qconfig(True)

其中,cali_qconfig_setterqat_qconfig_setter 分别为calibration模型和qat模型对应的qconfig设置,关于qconfig的设置方法与调试步骤,比如默认qconfig的设置、量化敏感算子设置等,请阅读 量化感知训练-Qconfig详解 章节的内 容。

训练策略不同

正如我们之前所说,量化训练其实是在纯浮点训练基础上的finetue。因此量化训练的时候,我们的初始学习率设置为浮点训练的十分之一,训练的epoch次数也大大减少,最重要的是 model 定义的时候,我们的 retrained 需要设置成已经训练出来的纯浮点模型的地址。

做完这些简单的调整之后,就可以开始训练我们的量化模型了。