FCOS3D检测模型训练

这篇教程以FCOS3D-efficientnetb0为例,告诉大家如何使用HAT算法包训练一个定点的3D检测模型。 在开始量化感知训练,也就是定点模型训练之前,首先需要训练一个精度较高的纯浮点模型,然后基于这个纯浮点模型做finetune,就可以快速的训练出定点模型。 所以我们从训练一个纯浮点的FCOS3D-efficientnetb0模型开始讲起。

训练流程

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

数据集准备

在开始训练模型之前,第一步是需要准备好数据集。这里我们训练 FCOS3D 模型使用的是开源的 nuscenes 数据集。解压缩之后数据目录结构如下所示:

tmp_data
|-- nuscenes
  |-- v1.0-mini.tar
  |-- v1.0-test_blobs.tar
  |-- v1.0-test_meta.tar
  |-- v1.0-trainval01_blobs.tar
  |-- ...
  |-- v1.0-trainval10_blobs.tar
  |-- v1.0-trainval_meta.tar
  |-- can_bus
  |-- maps
  |-- meta
  |-- samples
  |-- sweeps
  |-- v1.0-mini
  |-- v1.0-test
  |-- v1.0-trainval

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

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

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

tmp_data
|-- nuscenes
  |-- train_lmdb
  |-- val_lmdb
  |-- meta

train_lmdb和val_lmdb就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集。

模型训练

数据集准备好之后,就可以开始训练浮点型的FCOS3D-efficientnetb0检测网络了。 如果你只是单纯的想启动这样的训练任务,只需要运行下面的命令就可以:

python3 tools/train.py --stage "float" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

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

导出定点模型

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

python3 tools/export_hbir.py --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

模型验证

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

python3 tools/predict.py --stage "float" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py
python3 tools/predict.py --stage "calibration" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py
python3 tools/predict.py --stage "qat" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

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

python3 tools/predict.py --stage "int_infer" configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

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

仿真上板精度验证

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

python3 tools/validation_hbir.py --stage "align_bpu" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

模型推理和结果可视化

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

python3 tools/infer_hbir.py -c configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py --model-inputs ${model_inputs} --save-path ${save_path} --use-dataset

python3 tools/infer_hbir.py -c configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py --model-inputs ${model_inputs} --save-path ${save_path}

定点模型检查和编译

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

python3 tools/model_checker.py --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

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

python3 tools/compile_perf_hbir.py --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

训练细节

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

模型构建

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

model = dict(
    type="FCOS3D",
    backbone=dict(
        type="efficientnet",
        bn_kwargs=bn_kwargs,
        model_type="b0",
        num_classes=1000,
        include_top=False,
        activation="relu",
        use_se_block=False,
    ),
    neck=dict(
        type="BiFPN",
        in_strides=[2, 4, 8, 16, 32],
        out_strides=[8, 16, 32, 64, 128],
        stride2channels=dict({2: 16, 4: 24, 8: 40, 16: 112, 32: 320}),
        out_channels=64,
        num_outs=5,
        stack=3,
        start_level=2,
        end_level=-1,
        fpn_name="bifpn_sum",
    ),
    head=dict(
        type="FCOS3DHead",
        num_classes=10,
        in_channels=64,
        feat_channels=256,
        stacked_convs=2,
        strides=[8, 16, 32, 64, 128],
        group_reg_dims=(2, 1, 3, 1, 2),  # offset, depth, size, rot, velo
        use_direction_classifier=True,
        pred_attrs=True,
        num_attrs=9,
        cls_branch=(256,),
        reg_branch=(
            (256,),  # offset
            (256,),  # depth
            (256,),  # size
            (256,),  # rot
            (),  # velo
        ),
        dir_branch=(256,),
        attr_branch=(256,),
        centerness_branch=(64,),
        centerness_on_reg=True,
        return_for_compiler=False,
        output_int32=True,
    ),
    targets=dict(
        type="FCOS3DTarget",
        num_classes=10,
        background_label=None,
        bbox_code_size=9,
        regress_ranges=((-1, 48), (48, 96), (96, 192), (192, 384), (384, INF)),
        strides=[8, 16, 32, 64, 128],
        pred_attrs=True,
        num_attrs=9,
        center_sampling=True,
        center_sample_radius=1.5,
        centerness_alpha=2.5,
        norm_on_bbox=True,
    ),
    post_process=dict(
        type="FOCS3DPostProcess",
        num_classes=10,
        use_direction_classifier=True,
        strides=[8, 16, 32, 64, 128],
        group_reg_dims=(2, 1, 3, 1, 2),
        pred_attrs=True,
        num_attrs=9,
        attr_background_label=9,
        bbox_coder=dict(type="FCOS3DBBoxCoder", code_size=9),
        bbox_code_size=9,
        dir_offset=0.7854,
        test_cfg=dict(
            use_rotate_nms=True,
            nms_across_levels=False,
            nms_pre=1000,
            nms_thr=0.8,
            score_thr=0.05,
            min_bbox_size=0,
            max_per_img=100,
        ),
    ),
    loss=dict(
        type="FCOS3DLoss",
        num_classes=10,
        pred_attrs=True,
        group_reg_dims=(2, 1, 3, 1, 2),
        num_attrs=9,
        pred_velo=True,
        use_direction_classifier=True,
        dir_offset=0.7854,
        dir_limit_offset=0,
        diff_rad_by_sin=True,
        loss_cls=dict(
            type="FocalLoss",
            num_classes=11,
            gamma=2.0,
            alpha=0.25,
        ),
        loss_bbox=dict(type="SmoothL1Loss", beta=1.0 / 9.0, loss_weight=1.0),
        loss_dir=dict(
            type="CrossEntropyLoss", use_sigmoid=False, loss_weight=1.0
        ),
        loss_attr=dict(
            type="CrossEntropyLoss", use_sigmoid=False, loss_weight=1.0
        ),
        loss_centerness=dict(
            type="CrossEntropyLoss",
            use_sigmoid=True,
            loss_weight=1.0,
        ),
        train_cfg=dict(
            allowed_border=0,
            code_weight=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.05, 0.05],
            pos_weight=-1,
            debug=False,
        ),
    ),
)

其中, model 下面的 type 表示定义的模型名称,剩余的变量表示模型的其他组成部分。 这样定义模型的好处在于我们可以很方便的替换我们想要的结构。

数据增强

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

data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="NuscenesMonoDataset",
        data_path="./tmp_data/nuscenes/train_lmdb/",
        transforms=[
            dict(
                type="Pad",
                divisor=128,
            ),
            dict(
                type="ToTensor",
                to_yuv=True,
            ),
            dict(
                type="Normalize",
                mean=128.0,
                std=128.0,
            ),
        ],
    ),
    sampler=dict(type=torch.utils.data.DistributedSampler),
    batch_size=batch_size_per_gpu,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    collate_fn=hat.data.collates.collate_2d,
)

训练策略

为了训练一个精度高的模型,好的训练策略是必不可少的。对于每一个训练任务而言,相应的训练策略同样都定义在其中的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/efficientnet_imagenet/float-checkpoint-best.pth.tar"  # noqa: E501
                ),
                allow_miss=True,
                ignore_extra=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.SGD,
        params={"weight": dict(weight_decay=5e-5)},
        lr=0.001,
        momentum=0.9,
    ),
    batch_processor=batch_processor,
    num_epochs=12,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(
            type="StepDecayLrUpdater",
            warmup_len=0.3,
            lr_decay_id=[8, 11],
            step_log_interval=10,
        ),
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    sync_bn=True,
)

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

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

量化模型训练

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

python3 tools/train.py --stage "calibration" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py
python3 tools/train.py --stage "qat" --config configs/detection/fcos3d/fcos3d_efficientnetb0_nuscenes.py

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

calibration_trainer = dict(
    type="Calibrator",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        qconfig_params=dict(
            activation_calibration_observer="min_max",
        ),
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
            ),
            dict(type="Float2Calibration", convert_mode=convert_mode),
        ],
    ),
    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="NuscenesMonoMetric",
        data_root=meta_rootdir,
        version="v1.0-trainval",
    ),
    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),
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "calibration-checkpoint-best.pth.tar"
                ),
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        params={"weight": dict(weight_decay=0.01)},
        lr=1e-6,
    ),
    batch_processor=batch_processor,
    num_epochs=10,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="NuscenesMonoMetric",
        data_root=meta_rootdir,
        version="v1.0-trainval",
    ),
)

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 需要设置成已经训练出来的纯浮点模型的地址。

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