MapTROE 模型训练

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

训练流程

如果你只是想简单的把 MapTROE 的模型训练起来,那么可以首先阅读一下这一章的内容。和其他任务一样,对于所有的训练,评测任务,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

上面这两条命令分别对应着转换训练数据集和验证数据集。此外,MapTROE模型还需要导航地图(SD map)进行地图融合来提升模型精度,导航地图可以从开放街道地图(OpenStreetMap, OSM) https://www.openstreetmap.org 中获取。 打包完成之后,data目录下的文件结构应该如下所示:

tmp_data
    |-- nuscenes
        |-- meta
        |-- osm
        |-- v1.0-trainval 
            |-- train_lmdb
            |-- val_lmdb

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

模型训练

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

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

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

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

导出定点模型

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

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

模型验证

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

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

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

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

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

模型推理和结果可视化

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

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

python3 tools/infer_hbir.py --config configs/map/maptroe_henet_tinym_bevformer_nuscenes.py --model-inputs ${model_inputs} --save-path ${save_path}

定点模型检查和编译

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

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

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

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

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

训练细节

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

模型构建

model = dict(
    type="MapTROE",
    out_indices=(-1,),
    sd_map_fusion=True,
    backbone=dict(
        type="HENet",
        in_channels=3,
        block_nums=depth,
        embed_dims=width,
        attention_block_num=attention_block_num,
        mlp_ratios=mlp_ratios,
        mlp_ratio_attn=mlp_ratio_attn,
        act_layer=act_layer,
        use_layer_scale=use_layer_scale,
        layer_scale_init_value=1e-5,
        num_classes=1000,
        include_top=False,
        extra_act=extra_act,
        final_expand_channel=final_expand_channel,
        feature_mix_channel=feature_mix_channel,
        block_cls=block_cls,
        down_cls=down_cls,
        patch_embed=patch_embed,
        stage_out_norm=True,
    ),
    neck=dict(
        type="FPN",
        in_strides=[32],
        in_channels=[384],
        out_strides=[32],
        out_channels=[bev_embed_dims],
        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=bev_embed_dims,
        queue_length=queue_length,
        in_indices=(-1,),
        single_bev=single_bev,
        use_lidar2img=use_lidar2img,
        max_camoverlap_num=max_camoverlap_num,
        virtual_bev_h=int(bev_sparse_rate * bev_h_),
        virtual_bev_w=bev_w_,
        positional_encoding=dict(
            type="PositionEmbeddingLearned",
            num_pos_feats=bev_embed_dims // 2,
            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=bev_embed_dims,
            encoder_layer=dict(
                type="SingleBEVFormerEncoderLayer",
                embed_dims=bev_embed_dims,
                selfattention=dict(
                    type="HorizonMultiScaleDeformableAttention",
                    embed_dims=bev_embed_dims,
                    num_levels=1,
                    grid_align_num=10,
                    batch_first=True,
                    feats_size=[[bev_w_, bev_h_]],
                ),
                crossattention=dict(
                    type="HorizonSpatialCrossAttention",
                    max_camoverlap_num=max_camoverlap_num,
                    bev_h=bev_h_,
                    bev_w=bev_w_,
                    deformable_attention=dict(
                        type="HorizonMultiScaleDeformableAttention3D",
                        embed_dims=bev_embed_dims,
                        num_points=8,
                        num_levels=_num_levels_,
                        grid_align_num=20,
                        feats_size=[[25, 15]],
                    ),
                    embed_dims=bev_embed_dims,
                ),
                dropout=0.1,
            ),
        ),
    ),
    osm_encoder=dict(
        type="ConvDown",
        in_dim=1,
        mid_dim=hidden_dim // 2,
        out_dim=hidden_dim,
        quant_input=True,
    ),
    bev_fusion=dict(
        type="MapFusion",
        input_dim=bev_embed_dims,
        embed_dims=hidden_dim,
        bev_h=bev_h_,
        bev_w=bev_w_,
        bev_down=dict(
            type="ConvDown",
            in_dim=bev_embed_dims,
            mid_dim=bev_embed_dims,
            out_dim=bev_embed_dims,
            quant_input=False,
        ),
        fusion_up=dict(
            type="ConvUp",
            in_dim=2 * hidden_dim,
            mid_dim=bev_embed_dims,
            out_dim=bev_embed_dims,
        ),
    ),
    bev_decoders=[
        dict(
            type="MapInstanceDetectorHead",
            in_channels=bev_embed_dims,
            num_cam=6,
            bev_h=bev_h_,
            bev_w=bev_w_,
            embed_dims=head_embed_dims,
            num_vec_one2one=num_vec_one2one,
            num_vec_one2many=num_vec_one2many,
            k_one2many=6,
            num_pts_per_vec=fixed_ptsnum_per_pred_line,
            num_pts_per_gt_vec=fixed_ptsnum_per_gt_line,
            transform_method="minmax",
            gt_shift_pts_pattern="v2",
            code_size=2,
            num_classes=num_map_classes,
            aux_seg=aux_seg_cfg,
            decoder=dict(
                type="MapInstanceDecoder",
                num_layers=6,
                return_intermediate=True,
                decoder_layer=dict(
                    type="DetrTransformerDecoderLayer",
                    embed_dims=head_embed_dims,
                    crossattention=dict(
                        type="HorizonMultiPointDeformableAttention",
                        embed_dims=head_embed_dims,
                        num_levels=1,
                        grid_align_num=2,
                        num_points=fixed_ptsnum_per_pred_line,
                        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=4.0),
                    pts_cost=dict(
                        type="OrderedPtsL1Cost", weight=2.5, beta=0.01
                    ),
                    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=4.0,
                    reduction="mean",
                ),
                loss_pts=dict(type="PtsL1Loss", loss_weight=2.5, beta=0.01),
                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,
        sd_map_path=sd_map_path,
        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),
    shuffle=False,
    batch_size=batch_size_per_gpu,
    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,
        sd_map_path=sd_map_path,
        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 表示地图数据的路径,sd_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=(pretrained_model_dir),
                allow_miss=True,
                ignore_extra=True,
                verbose=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        params={
            "backbone": dict(lr_mult=0.1),
        },
        lr=float_lr,
        weight_decay=0.1,
    ),
    batch_processor=batch_processor,
    device=None,
    num_epochs=30,
    callbacks=[
        stat_callback,
        loss_show_update,
        grad_callback,
        dict(
            type="CosineAnnealingLrUpdater",
            warmup_len=500,
            warmup_by="step",
            warmup_lr_ratio=1.0 / 3,
            warmup_lr_begin2ratio=True,
            step_log_interval=500,
            stop_lr=3e-3 * float_lr,
        ),
        ckpt_callback,
        val_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=calib_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="RepModel2Deploy",
            ),
            dict(
                type="Float2Calibration",
                convert_mode=convert_mode,
                example_data_loader=calibration_example_data_loader,
                qconfig_setter=cali_qconfig_setter,
            ),
            dict(
                type="FixWeightQScale",
            ),
        ],
    ),
    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="RepModel2Deploy",
            ),
            dict(
                type="Float2QAT",
                convert_mode=convert_mode,
                example_data_loader=copy.deepcopy(
                    calibration_example_data_loader
                ),
                qconfig_setter=qat_qconfig_setter,
            ),
            dict(
                type="FixWeightQScale",
            ),
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "calibration-checkpoint-best.pth.tar"
                ),
                allow_miss=True,
                ignore_extra=True,
                verbose=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        params={
            "backbone": dict(lr_mult=0.1),
        },
        lr=qat_lr,
        weight_decay=0.1,
    ),
    batch_processor=qat_batch_processor,
    device=None,
    num_epochs=3,
    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设置如下:

q_templates = [
    ModuleNameTemplate({"": qint8}),
    MatmulDtypeTemplate(
        input_dtypes=[qint8, qint8],
    ),
    ConvDtypeTemplate(
        input_dtype=qint8,
        weight_dtype=qint8,
    ),
]
cali_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(
        observer=(observer_v2.MSEObserver)
    ),
    templates=q_templates,
    enable_optimize=True,
    save_dir=ckpt_dir,
    custom_qconfig_mapping=None,
)
qat_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(
        observer=(observer_v2.MinMaxObserver), fix_scale=True
    ),
    templates=q_templates,
    enable_optimize=True,
    save_dir=ckpt_dir,
    custom_qconfig_mapping=None,
)

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