lidarMultiTask模型训练

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

数据集准备

在开始训练模型之前,第一步是需要准备好数据集,可以在 nuscenes 数据集 下载 v1.0 版本的完整数据文件和 nuScenes-lidarseg 的完整数据文件。

下载后,解压并按照如下方式组织文件夹结构,其中,lidarseg文件夹摆放规则可参照 nuscenes 官方Tutorials

├── data │ ├── nuscenes │ │ ├── lidarseg │ │ ├── maps │ │ ├── samples │ │ ├── sweeps │ │ ├── v1.0-mini │ │ ├── v1.0-trainval

为了提升训练的速度,我们对数据信息文件做了一个打包,将其转换成lmdb格式的数据集。其中,lidarMultiTask 模型只使用了 nuscenes 数据集的点云部分文件。 只需要运行下面的脚本,就可以成功实现转换:

python3 tools/datasets/nuscenes_packer.py --src-data-dir data/nuscenes/ --pack-type lmdb --target-data-dir tmp_data/nuscenes/lidar_seg/v1.0-trainval --version v1.0-trainval --split-name val --only-lidar python3 tools/datasets/nuscenes_packer.py --src-data-dir data/nuscenes/ --pack-type lmdb --target-data-dir tmp_data/nuscenes/lidar_seg/v1.0-trainval --version v1.0-trainval --split-name train --only-lidar

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

├── tmp_data │ ├── nuscenes │ │ ├── lidar_seg │ │ │ ├── v1.0-trainval │ │ │ │ ├── train_lmdb # 新生成的 lmdb │ │ │ │ ├── val_lmdb # 新生成的 lmdb │ │ ├── meta │ │ │ ├── maps │ │ │ ├── v1.0-mini │ │ │ ├── v1.0-trainval

train_lmdbval_lmdb 就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集,meta中为评测脚本需要的初始化信息,具体信息是从nuscenes原始数据集中拷贝得来。

同时,为了训练nuscenes点云数据,还需要为nuscenes数据集生成每个单独的训练目标的点云数据,并将其存储在 tmp_nuscenes/lidar/nuscenes_gt_database.bin 格式的文件中,文件保存目录可根据需要更改。同时,需要为这部分数据生成 .pkl 格式的包含数据信息的文件。 此外,训练数据集初始化过程中需要读取全部数据集中每个sample的类别信息,并进行重采样操作,我们可以提前生成对应的信息并保存成 .pkl 格式的文件,可以加速训练过程。 通过运行下面的命令来创建上述数据:

python3 tools/create_data.py --dataset nuscenes --root-dir ./tmp_data/nuscenes/lidar_seg/v1.0-trainval --extra-tag nuscenes --out-dir tmp_nuscenes/lidar

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

├── tmp_data │ ├── nuscenes │ │ ├── lidar_seg │ │ │ ├── v1.0-trainval │ │ │ │ ├── train_lmdb # 新生成的 lmdb │ │ │ │ ├── val_lmdb # 新生成的 lmdb │ │ ├── meta │ │ │ ├── maps │ │ │ ├── v1.0-mini │ │ │ ├── v1.0-trainval ├── tmp_nuscenes │ ├── lidar │ │ ├── nuscenes_gt_database # 新生成的 nuscenes_gt_database │ │ │ ├── xxxxx.bin │ │ ├── nuscenes_dbinfos_train.pkl # 新生成的 nuscenes_dbinfos_train.pkl │ │ ├── nuscenes_infos_train.pkl # 新生成的 nuscenes_infos_train.pkl

其中, nuscenes_gt_databasenuscenes_dbinfos_train.pkl 是训练是用于采样的样本,而 nuscenes_dbinfos_train.pkl 则包含训练数据集初始化需要的信息。

浮点模型训练

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

python3 tools/calops.py --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py --method hook

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

python3 tools/train.py --stage float --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

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

模型构建

lidarMultiTask 的网络结构主要借鉴 CenterPoint 模型,在模型层面,主要修改了neck部分并添加了分割的输出头,详细可参考下面的config文件。 我们通过在config配置文件中定义 model 这样的一个dict型变量,就可以方便的实现对模型的定义和修改。

model = dict( type="LidarMultiTask", feature_map_shape=get_feature_map_size(point_cloud_range, voxel_size), pre_process=dict( type="CenterPointPreProcess", pc_range=point_cloud_range, voxel_size=voxel_size, max_voxels_num=max_voxels, max_points_in_voxel=max_num_points, norm_range=[-51.2, -51.2, -5.0, 0.0, 51.2, 51.2, 3.0, 255.0], norm_dims=[0, 1, 2, 3], ), reader=dict( type="PillarFeatureNet", num_input_features=5, num_filters=(64,), with_distance=False, pool_size=(max_num_points, 1), voxel_size=voxel_size, pc_range=point_cloud_range, bn_kwargs=norm_cfg, quantize=True, use_4dim=True, use_conv=True, hw_reverse=True, ), scatter=dict( type="PointPillarScatter", num_input_features=64, use_horizon_pillar_scatter=True, quantize=True, ), backbone=dict( type="MixVarGENet", net_config=net_config, disable_quanti_input=True, input_channels=64, input_sequence_length=1, num_classes=1000, bn_kwargs=bn_kwargs, include_top=False, bias=True, output_list=[0, 1, 2, 3, 4], ), neck=dict( type="Unet", in_strides=(2, 4, 8, 16, 32), out_strides=(4,), stride2channels=dict( { 2: 64, 4: 64, 8: 64, 16: 96, 32: 160, } ), out_stride2channels=dict( { 2: 128, 4: 128, 8: 128, 16: 128, 32: 160, } ), factor=2, group_base=8, bn_kwargs=bn_kwargs, ), lidar_decoders=[ dict( type="LidarSegDecoder", name="seg", task_weight=80.0, task_feat_index=0, head=dict( type="DepthwiseSeparableFCNHead", input_index=0, in_channels=128, feat_channels=64, num_classes=2, dropout_ratio=0.1, num_convs=2, bn_kwargs=bn_kwargs, int8_output=False, ), target=dict( type="FCNTarget", ), loss=dict( type="CrossEntropyLoss", loss_name="seg", reduction="mean", ignore_index=-1, use_sigmoid=False, class_weight=[1.0, 10.0], ), decoder=dict( type="FCNDecoder", upsample_output_scale=4, use_bce=False, bg_cls=-1, ), ), dict( type="LidarDetDecoder", name="det", task_weight=1.0, task_feat_index=0, head=dict( type="DepthwiseSeparableCenterPointHead", in_channels=128, tasks=tasks, share_conv_channels=64, share_conv_num=1, common_heads=common_heads, head_conv_channels=64, init_bias=-2.19, final_kernel=3, ), target=dict( type="CenterPointLidarTarget", grid_size=[512, 512, 1], voxel_size=voxel_size, point_cloud_range=point_cloud_range, tasks=tasks, dense_reg=1, max_objs=500, gaussian_overlap=0.1, min_radius=2, out_size_factor=4, norm_bbox=True, with_velocity=with_velocity, ), loss=dict( type="CenterPointLoss", loss_cls=dict(type="GaussianFocalLoss", loss_weight=1.0), loss_bbox=dict( type="L1Loss", reduction="mean", loss_weight=0.25, ), with_velocity=with_velocity, code_weights=[ 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.2, 0.2, ], ), decoder=dict( type="CenterPointPostProcess", tasks=tasks, norm_bbox=True, bbox_coder=dict( type="CenterPointBBoxCoder", pc_range=point_cloud_range[:2], post_center_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0], max_num=100, score_threshold=0.1, out_size_factor=4, voxel_size=voxel_size[:2], ), # test_cfg max_pool_nms=False, score_threshold=0.1, post_center_limit_range=[ -61.2, -61.2, -10.0, 61.2, 61.2, 10.0, ], min_radius=[4, 12, 10, 1, 0.85, 0.175], out_size_factor=4, nms_type="rotate", pre_max_size=1000, post_max_size=83, nms_thr=0.2, box_size=9, ), ), ], )

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

数据增强

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

train_dataset = dict( type="NuscenesLidarWithSegDataset", num_sweeps=9, data_path=os.path.join(data_rootdir, "train_lmdb"), info_path=os.path.join(gt_data_root, "nuscenes_infos_train.pkl"), load_dim=5, use_dim=[0, 1, 2, 3, 4], pad_empty_sweeps=True, remove_close=True, use_valid_flag=True, classes=det_class_names, transforms=[ dict( type="LidarMultiPreprocess", class_names=det_class_names, global_rot_noise=[-0.3925, 0.3925], global_scale_noise=[0.95, 1.05], db_sampler=db_sampler, ), dict( type="ObjectRangeFilter", point_cloud_range=point_cloud_range, ), dict( type="AssignSegLabel", bev_size=[512, 512], num_classes=2, class_names=[0, 1], point_cloud_range=point_cloud_range, voxel_size=voxel_size[:2], ), dict(type="LidarReformat", with_gt=True), ], ) data_loader = dict( type=torch.utils.data.DataLoader, dataset=dict(type="CBGSDataset", dataset=train_dataset), sampler=dict(type=torch.utils.data.DistributedSampler), batch_size=batch_size_per_gpu, shuffle=False, num_workers=4, pin_memory=False, collate_fn=hat.data.collates.collate_lidar3d, )

其中type直接用的pytorch自带的接口 torch.utils.data.DataLoader,表示的是将 batch_size 大小的样本组合到一起。 这里面唯一需要关注的可能是 dataset 这个变量,data_path 路径也就是我们在第一部分数据集准备中提到的路径。transforms 下面包含着一系列的数据增强。而 val_data_loader 中只有读取点云文件、生成分割标签、数据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=2e-4, weight_decay=0.01, ), batch_processor=batch_processor, num_epochs=20, 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=200, ), grad_callback, val_callback, ckpt_callback, ], sync_bn=True, train_metrics=dict( type="LossShow", ), )

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

注解

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

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

量化模型训练

当我们有了纯浮点模型之后,就可以开始训练相应的定点模型了。和浮点训练的方式一样,我们只需要通过运行下面的脚本就可以训练定点模型了。 不过这里需要说明的是,lidarMultiTask模型在量化训练过程中建议加上calibration的流程。calibration可以为QAT的量化训练提供一个更好的初始化参数。

python3 tools/train.py --stage calibration --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py python3 tools/train.py --stage qat --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

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

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), dict( type="LoadCheckpoint", checkpoint_path=os.path.join( ckpt_dir, "calibration-checkpoint-last.pth.tar" ), ), ], ), data_loader=data_loader, optimizer=dict( type=torch.optim.SGD, weight_decay=0.0, lr=1e-4, momentum=0.9, ), batch_processor=batch_processor, num_epochs=10, 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=200, ), grad_callback, val_callback, ckpt_callback, ], train_metrics=dict( type="LossShow", ), )

model_convert_pipeline参数值不同

当我们训练量化模型的时候,需要设置model_convert_pipeline,此时相应的浮点模型会被转换成量化模型,相关代码如下:

model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode="fuse_bn", converters=[ dict(type="Float2QAT", convert_mode=convert_mode), dict( type="LoadCheckpoint", checkpoint_path=os.path.join( ckpt_dir, "qat-checkpoint-best.pth.tar" ), ), dict(type="QAT2Quantize", convert_mode=convert_mode), ], )

关于量化训练中的关键步骤,比如准备浮点模型、算子替换、插入量化和反量化节点、设置量化参数以及算子的融合等,请阅读 量化感知训练 章节的内容。

训练策略不同

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

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

导出定点模型

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

python3 tools/export_hbir.py --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

模型验证

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

python3 tools/predict.py --stage float --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py python3 tools/predict.py --stage qat --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

同时,我们还提供了quantization模型的性能测试,只需要运行以下命令,但需要注意是必须要先导出hbir:

python3 tools/predict.py --stage "int_infer" --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

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

仿真上板精度验证

除了上述模型验证之外,我们还提供和上板完全一致的精度验证方法,可以通过下面的方式完成:

python3 tools/validation_hbir.py --stage "align_bpu" --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py

模型推理和结果可视化

如果你希望可以看到训练出来的模型对于雷达点云的检测和分割效果,我们的tools文件夹下面同样提供了点云预测及可视化的脚本,你只需要运行以下脚本即可:

python3 tools/infer_hbir.py --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py --model-inputs input_points:${lidar-pointcloud-path} --save-path ${save_path}

模型检查和编译

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

python3 tools/compile_perf_hbir.py --config configs/lidar_multi_task/centerpoint_mixvargnet_multitask_nuscenes.py