MapTR 模型训练

MapTR 参考算法基于 Horizon Algorithm Toolkit(HAT,地平线自研深度学习算法包)开发。训练config位于 configs/map/ 路径下。 下文以configs/map/maptrv2_resnet50_bevformer_nuscenes.py为例介绍如何配置并训练 MapTR 参考算法。

训练流程

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

数据集准备

这里以nuscense数据集为例,可以从 https://www.nuscenes.org/nuscenes 下载数据集 。同时,为了提升训练的速度,我们对原始的jpg格式的数据集做了一个打包,将其转换成lmdb格式的数据集。只需要运行下面的脚本,就可以成功实现转换:

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

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

tmp_data |-- nuscenes |-- metas |-- v1.0-trainval |-- train_lmdb |-- val_lmdb

train_lmdb和val_lmdb就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集。metas中包含模型需要的地图信息。

模型训练

数据集准备好之后,就可以开始 MapTR 模型的训练了。

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

python3 tools/train.py --stage float --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py python3 tools/train.py --stage calibration --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py python3 tools/train.py --stage qat --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

以上命令分别完成float、calibration、qat阶段的训练,其中calibration阶段的训练需要以训练好的浮点模型为基础,qat阶段的训练则需要训练好的calibration模型为基础,具体内容请阅读 量化感知训练 章节的内容。

导出定点模型

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

python3 tools/export_hbir.py --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

模型验证

在完成训练之后,可以得到训练完成的float、calibration和qat模型。和训练方法类似,我们可以用相同方法来对训好的模型做指标验证,得到为 FloatCalibrationQAT 的指标。

python3 tools/predict.py --stage float --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py python3 tools/predict.py --stage calibration --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py python3 tools/predict.py --stage qat --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

和训练模型时类似,--stage 后面的参数为 floatcalibrationqat 时,分别可以完成对训练好的浮点模型、calibration模型、qat模型的验证。

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

python3 tools/predict.py --stage int_infer --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

模型推理

HAT 提供了 infer_hbir.py 脚本提供了对定点模型的推理结果进行可视化展示:

python3 tools/infer_hbir.py --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py --model-inputs img:${img-path} --save-path ${save_path}

仿真上板精度验证

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

python3 tools/validation_hbir.py --stage align_bpu --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

定点模型检查和编译

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

python3 tools/model_checker.py --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

在模型训练完成后,可以通过 compile_perf_hbir 脚本将量化模型编译成可以上板运行的 hbm 文件:

python3 tools/compile_perf_hbir.py --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

以上就是从数据准备到生成量化可部署模型的全过程。

训练细节

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

模型构建

model = dict( type="MapTRv2", out_indices=(-1,), backbone=dict( type="ResNet50", num_classes=1000, bn_kwargs=bn_kwargs, include_top=False, ), neck=dict( type="FPN", in_strides=[32], in_channels=[2048], out_strides=[32], out_channels=[_dim_], bn_kwargs=dict(eps=1e-5, momentum=0.1), ), view_transformer=dict( type="SingleBevFormerViewTransformer", bev_h=bev_h_, bev_w=bev_w_, pc_range=point_cloud_range, num_points_in_pillar=4, embed_dims=_dim_, queue_length=queue_length, in_indices=(-1,), single_bev=single_bev, use_lidar2img=use_lidar2img, positional_encoding=dict( type="LearnedPositionalEncoding", num_feats=_pos_dim_, row_num_embed=bev_h_, col_num_embed=bev_w_, ), encoder=dict( type="SingleBEVFormerEncoder", num_layers=1, return_intermediate=False, bev_h=bev_h_, bev_w=bev_w_, embed_dims=_dim_, encoder_layer=dict( type="SingleBEVFormerEncoderLayer", embed_dims=_dim_, selfattention=dict( type="HorizonMSDeformableAttention", embed_dims=_dim_, num_levels=1, view_gird_in=4, view_gird_out=4, batch_first=True, feats_size=[[bev_w_, bev_h_]], ), crossattention=dict( type="HorizonSpatialCrossAttention", view_num=8, deformable_attention=dict( type="HorizonMSDeformableAttention3D", embed_dims=_dim_, num_points=8, num_levels=_num_levels_, view_gird_in=160, view_gird_out=100, feats_size=[[25, 15]], ), embed_dims=_dim_, ), dropout=0.1, ), ), ), bev_decoders=[ dict( type="MapTRPerceptionDecoderv2", bev_h=bev_h_, bev_w=bev_w_, embed_dims=_dim_, pc_range=point_cloud_range, queue_length=queue_length, numcam=6, num_vec_one2one=50, num_vec_one2many=300, k_one2many=6, num_pts_per_vec=fixed_ptsnum_per_pred_line, num_pts_per_gt_vec=fixed_ptsnum_per_gt_line, dir_interval=1, query_embed_type="instance_pts", transform_method="minmax", gt_shift_pts_pattern="v2", num_classes=num_map_classes, code_size=2, aux_seg=aux_seg_cfg, decoder=dict( type="MapTRDecoder", num_layers=6, return_intermediate=True, decoder_layer=dict( type="DecoupledDetrTransformerDecoderLayer", embed_dims=_dim_, crossattention=dict( type="HorizonMSDeformableAttention", embed_dims=_dim_, num_levels=1, view_gird_out=4, view_gird_in=4, feats_size=[[bev_w_, bev_h_]], ), dropout=0.1, ), ), criterion=dict( type="MapTRCriterion", dir_interval=1, num_classes=num_map_classes, code_weights=[1.0, 1.0, 1.0, 1.0], sync_cls_avg_factor=True, pc_range=point_cloud_range, num_pts_per_vec=fixed_ptsnum_per_pred_line, # one bbox num_pts_per_gt_vec=fixed_ptsnum_per_gt_line, gt_shift_pts_pattern="v2", aux_seg=aux_seg_cfg, assigner=dict( type="MapTRAssigner", cls_cost=dict(type="FocalLossCost", weight=2.0), pts_cost=dict(type="OrderedPtsL1Cost", weight=5), pc_range=point_cloud_range, ), loss_cls=dict( type="FocalLoss", loss_name="cls", num_classes=num_map_classes + 1, alpha=0.25, gamma=2.0, loss_weight=2.0, reduction="mean", ), loss_pts=dict(type="PtsL1Loss", loss_weight=5.0), loss_dir=dict(type="PtsDirCosLoss", loss_weight=0.005), loss_seg=dict( type="SimpleLoss", pos_weight=4.0, loss_weight=1.0 ), loss_pv_seg=dict( type="SimpleLoss", pos_weight=1.0, loss_weight=2.0 ), ), post_process=dict( type="MapTRPostProcess", post_center_range=post_center_range, pc_range=point_cloud_range, max_num=50, num_classes=num_map_classes, ), ), ], )

其中,model 下面的 type 表示定义的模型名称,剩余的变量表示模型的其他组成部分。这样定义模型的好处在于我们可以很方便的替换我们想要的结构。例如,如果我们想训练一个backbone为resnet18的模型,只需要将 model 下面的 backbone 替换掉就可以。

数据增强

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

data_loader = dict( type=torch.utils.data.DataLoader, dataset=dict( type="NuscenesMapDataset", data_path=os.path.join(data_rootdir, "train_lmdb"), map_path=meta_rootdir, pc_range=point_cloud_range, test_mode=False, bev_size=(bev_h_, bev_w_), fixed_ptsnum_per_line=fixed_ptsnum_per_gt_line, padding_value=-10000, map_classes=map_classes, queue_length=queue_length, aux_seg=aux_seg_cfg, with_bev_bboxes=False, with_ego_bboxes=False, with_bev_mask=False, use_lidar_gt=use_lidar_gt, transforms=[ dict(type="MultiViewsImgResize", size=(450, 800)), dict( type="MultiViewsImgTransformWrapper", transforms=[ dict( type="TorchVisionAdapter", interface="ColorJitter", brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, ), dict(type="PILToNumpy"), dict( type="GridMask", use_h=True, use_w=True, rotate=1, offset=False, ratio=0.5, mode=1, prob=0.7, ), dict(type="ToTensor", to_yuv=False), dict(type="Pad", divisor=32), dict(type="BgrToYuv444", rgb_input=True), dict(type="Normalize", mean=128.0, std=128.0), ], ), ], ), sampler=dict(type=torch.utils.data.DistributedSampler), batch_size=batch_size_per_gpu, shuffle=False, num_workers=2, pin_memory=True, collate_fn=collate_nuscenes_sequencev2, ) val_data_loader = dict( type=torch.utils.data.DataLoader, dataset=dict( type="NuscenesMapDataset", data_path=os.path.join(data_rootdir, "val_lmdb"), map_path=meta_rootdir, pc_range=point_cloud_range, test_mode=True, bev_size=(bev_h_, bev_w_), fixed_ptsnum_per_line=fixed_ptsnum_per_gt_line, padding_value=-10000, map_classes=map_classes, queue_length=test_queue_length, with_bev_bboxes=False, with_ego_bboxes=False, with_bev_mask=False, use_lidar_gt=use_lidar_gt, transforms=[ dict(type="MultiViewsImgResize", size=(450, 800)), dict( type="MultiViewsImgTransformWrapper", transforms=[ dict(type="PILToTensor"), dict(type="Pad", divisor=32), dict(type="BgrToYuv444", rgb_input=True), dict(type="Normalize", mean=128.0, std=128.0), ], ), ], ), sampler=None, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, collate_fn=collate_nuscenes_sequencev2, )

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

训练策略

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

float_trainer = dict( type="distributed_data_parallel_trainer", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", converters=[ dict( type="LoadCheckpoint", checkpoint_path=( "./tmp_pretrained_models/resnet50_imagenet/float-checkpoint-best.pth.tar" # noqa: E501 ), allow_miss=True, ignore_extra=True, verbose=False, ), ], ), data_loader=data_loader, optimizer=dict( type=torch.optim.AdamW, params={ "backbone": dict(lr_mult=0.1), }, lr=float_lr, weight_decay=0.01, ), batch_processor=batch_processor, device=None, num_epochs=24, callbacks=[ stat_callback, loss_show_update, grad_callback, dict( type="CosineAnnealingLrUpdater", warmup_len=500, warmup_by="step", warmup_lr_ratio=1.0 / 3, step_log_interval=500, stop_lr=1e-3 * float_lr, ), val_callback, ckpt_callback, ], sync_bn=True, train_metrics=dict( type="LossShow", ), val_metrics=[ val_map_metric, ], )

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

注意

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

量化训练

使用 float_trainer 的设置,可以训练出来一个高精度的浮点模型。当我们有了浮点模型之后,就可以开始训练相应的量化模型了。此时我们使用的训练策略来自于config文件中的 calibration_trainerqat_trainer

calibration_trainer = dict( type="Calibrator", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode=qat_mode, converters=[ dict( type="LoadCheckpoint", checkpoint_path=os.path.join( ckpt_dir, "float-checkpoint-best.pth.tar" ), ignore_extra=True, verbose=True, check_hash=False, ), dict( type="Float2Calibration", convert_mode=convert_mode, example_data_loader=calibration_example_data_loader, qconfig_setter=cali_qconfig_setter, ), ], ), data_loader=calibration_data_loader, batch_processor=calibration_batch_processor, num_steps=calibration_step, device=None, callbacks=[ stat_callback, calibration_ckpt_callback, calibration_val_callback, ], val_metrics=[ val_map_metric, ], log_interval=calibration_step / 10, ) qat_trainer = dict( type="distributed_data_parallel_trainer", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode=qat_mode, 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" ), ignore_extra=True, verbose=True, ), ], ), data_loader=data_loader, optimizer=dict( type=torch.optim.AdamW, lr=qat_lr, weight_decay=0.01, ), batch_processor=qat_batch_processor, device=None, num_epochs=2, callbacks=[ stat_callback, loss_show_update, grad_callback, qat_val_callback, qat_ckpt_callback, ], sync_bn=True, train_metrics=dict( type="LossShow", ), val_metrics=[ val_map_metric, ], )

量化训练其实是在纯浮点训练基础上的finetue,因此量化训练的时候,学习率 qat_lr 要比 float_lr 小得多,qat训练的epoch次数也大大减少,最重要的是 model 定义的时候,calibration模型的 pretrained 需要设置成已经训练出来的纯浮点模型的地址,qat模型则需要设置为校准出来的calibration模型地址。

qconfig的设置

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

cali_qconfig_setter = ( sensitive_op_calibration_8bit_weight_16bit_act_qconfig_setter( pts_table, topk=30, ratio=None, ), default_calibration_qconfig_setter, ) qat_qconfig_setter = ( sensitive_op_qat_8bit_weight_16bit_fixed_act_qconfig_setter( pts_table, topk=30, ratio=None, ), default_qat_fixed_act_qconfig_setter, )

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

量化敏感度算子排序

量化训练过程中,我们需要将某些量化敏感的算子设置为int16,以满足模型的量化精度需求。量化敏感算子的排序可以通过运行以下命令获得。

python3 tools/quant_analysis.py --config configs/map/maptrv2_resnet50_bevformer_nuscenes.py

其中的关键步骤,请阅读 量化感知训练-精度调优工具使用指南 章节的内容。