10.4.16. CenterPoint检测模型训练¶
这篇教程主要是告诉大家如何利用HAT在数据集 nuscenes
上从头开始训练一个 CenterPoint
模型,包括浮点、量化和定点模型。
10.4.16.1. 数据集准备¶
在开始训练模型之前,第一步是需要准备好数据集,可以在 nuscenes 数据集 下载 1.0 版本的完整数据文件。
下载后,解压并按照如下方式组织文件夹结构:
├── data
│ ├── nuscenes
│ │ ├── maps
│ │ ├── samples
│ │ ├── sweeps
│ │ ├── v1.0-trainval
为了提升训练的速度,我们对数据信息文件做了一个打包,将其转换成lmdb格式的数据集。 只需要运行下面的脚本,就可以成功实现转换:
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-trainval
train_lmdb
和 val_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-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_database
和 nuscenes_dbinfos_train.pkl
是训练是用于采样的样本,而 nuscenes_dbinfos_train.pkl
则包含训练数据集初始化需要的信息。
10.4.16.2. 浮点模型训练¶
数据集准备好之后,就可以开始训练浮点型的CenterPoint网络了。在网络训练开始之前,你可以使用以下命令先测试一下网络的计算量和参数数量:
python3 tools/calops.py --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
如果你只是单纯的想启动这样的训练任务,只需要运行下面的命令就可以:
python3 tools/train.py --stage float --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
由于HAT算法包使用了一种巧妙的注册机制,使得每一个训练任务都可以按照这种train.py加上config配置文件的形式启动。 train.py是统一的训练脚本,与任务无关,我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的config配置文件里面。 config文件里面提供了模型构建、数据读取等关键的dict。
10.4.16.2.1. 模型构建¶
CenterPoint 的网络结构可以参考 论文 ,这里不做详细介绍。
我们通过在config配置文件中定义 model
这样的一个dict型变量,就可以方便的实现对模型的定义和修改。
model = dict(
type="CenterPointDetector",
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,
),
backbone=dict(
type="PointPillarScatter",
num_input_features=64,
use_horizon_pillar_scatter=True,
quantize=True,
),
neck=dict(
type="SECONDNeck",
in_feature_channel=64,
down_layer_nums=[3, 5, 5],
down_layer_strides=[2, 2, 2],
down_layer_channels=[64, 128, 256],
up_layer_strides=[0.5, 1, 2],
up_layer_channels=[128, 128, 128],
bn_kwargs=norm_cfg,
quantize=True,
use_relu6=False,
),
head=dict(
type="CenterPointHead",
in_channels=sum([128, 128, 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,
),
targets=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],
),
postprocess=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。
10.4.16.2.2. 数据增强¶
跟 model
的定义一样,数据增强的流程是通过在config配置文件中定义 data_loader
和 val_data_loader
这两个dict来实现的,
分别对应着训练集和验证集的处理流程。以 data_loader
为例:
train_dataset = dict(
type="NuscenesLidarDataset",
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=class_names,
transforms=[
dict(
type="PointCloudPreprocess",
mode="train",
current_mode="train",
class_names=class_names,
shuffle_points=True,
min_points_in_gt=-1,
flip_both=True,
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="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
中只有读取点云文件、Format 和 Collect3D。
你也可以通过在 transforms
中插入新的dict实现自己希望的数据增强操作。
10.4.16.2.3. 训练策略¶
为了训练一个精度高的模型,好的训练策略是必不可少的。对于每一个训练任务而言,相应的训练策略同样都定义在其中的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",
),
val_metrics=[val_nuscenes_metric],
)
float_trainer
从大局上定义了我们的训练方式,包括使用多卡分布式训练(distributed_data_parallel_trainer),模型训练的epoch次数,以及优化器的选择。
同时 callbacks
中体现了模型在训练过程中使用到的小策略以及用户想实现的操作,包括学习率的变换方式(CyclicLrUpdater),
在训练过程中验证模型的指标(Validation),以及保存(Checkpoint)模型的操作。当然,如果你有自己希望模型在训练过程中实现的操作,也可以按照这种dict的方式添加。
float_trainer
负责将整个训练的逻辑给串联起来,其中也会负责模型的pretrain。
注解
如果需要复现精度,config中的训练策略最好不要修改。否则可能会有意外的训练情况出现。
通过上面的介绍,你应该对config文件的功能有了一个比较清楚的认识。然后通过前面提到的训练脚本,就可以训练一个高精度的纯浮点的检测模型。 当然训练一个好的检测模型不是我们最终的目的,它只是做为一个pretrain为我们后面训练定点模型服务的。
10.4.16.3. 量化模型训练¶
当我们有了纯浮点模型之后,就可以开始训练相应的定点模型了。和浮点训练的方式一样,我们只需要通过运行下面的脚本就可以训练定点模型了。不过这里需要说明的是,CenterPoint在量化训练过程中建议加上calibration的流程。calibration可以为QAT的量化训练提供一个更好的初始化参数。
python3 tools/train.py --stage calibration --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
python3 tools/train.py --stage qat --config configs/detection/centerpoint/centerpoint_pointpillar_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=2e-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",
),
val_metrics=[val_nuscenes_metric],
)
10.4.16.3.1. quantize参数的值不同¶
当我们训练量化模型的时候,需要设置quantize=True,此时相应的浮点模型会被转换成量化模型,相关代码如下:
model.fuse_model()
model.set_qconfig()
horizon.quantization.prepare_qat(model, inplace=True)
关于量化训练中的关键步骤,比如准备浮点模型、算子替换、插入量化和反量化节点、设置量化参数以及算子的融合等, 请阅读 量化感知训练 章节的内容。
10.4.16.3.2. 训练策略不同¶
正如我们之前所说,量化训练其实是在纯浮点训练基础上的finetue。因此量化训练的时候,我们的初始学习率设置为浮点训练的十分之一,
训练的epoch次数也大大减少,最重要的是 model
定义的时候,我们的 pretrained
需要设置成已经训练出来的纯浮点模型的地址。
做完这些简单的调整之后,就可以开始训练我们的量化模型了。
10.4.16.3.3. 模型验证¶
模型训练完成之后,我们还可以验证训练出来的模型性能。由于我们提供了float和qat两阶段的训练过程,相应的我们可以验证这两个阶段训练出来的模型性能, 只需要相应的运行以下两条命令即可:
python3 tools/predict.py --stage float --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
python3 tools/predict.py --stage qat --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
同时,我们还提供了quantization模型的性能测试,只需要运行以下命令:
python3 tools/predict.py --stage int_infer --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
这个显示出来的精度才是最终的int8模型的真正精度,当然这个精度和qat验证阶段的精度应该是保持十分接近的。
10.4.16.3.4. 仿真上板精度验证¶
除了上述模型验证之外,我们还提供和上板完全一致的精度验证方法,可以通过下面的方式完成:
python3 tools/align_bpu_validation.py --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py
10.4.16.3.5. 结果可视化¶
如果你希望可以看到训练出来的模型对于雷达点云的检测效果,我们的tools文件夹下面同样提供了点云预测及可视化的脚本,你只需要运行以下脚本即可:
python3 tools/infer.py --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py --model-inputs input_points:${lidar-pointcloud-path} --save-path ${save_path}
10.4.16.3.6. 模型检查和编译¶
在训练完成之后,可以使用 compile
的工具用来将量化模型编译成可以上板运行的 hbm
文件,同时该工具也能预估在BPU上的运行性能,可以采用以下脚本:
python3 tools/compile_perf.py --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py --out-dir ./ --opt 3
10.4.16.3.7. ONNX模型导出¶
如果想要导出onnx模型, 运行下面的命令即可:
python3 tools/export_onnx.py --config configs/detection/centerpoint/centerpoint_pointpillar_nuscenes.py