• 中文
  • 3.8.1

    • Occupancy预测模型

    Occupancy预测模型基于Horizon Torch Samples（地平线自研深度学习框架）开发，关于Horizon Torch Samples的使用介绍可以参考Horizon Torch Samples使用文档。Occupancy预测模型的训练config位于configs/occ/路径下。 下文以configs/occ/flashocc_henet_lss_occ3d_nuscenes.py为例介绍如何配置并训练Occupancy预测模型。

    训练流程

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

    数据集准备

    这里以nuscense数据集为例，可以从 https://www.nuscenes.org/nuscenes 下载数据集, 对于Occupancy预测任务，还需要下载OCC的GT，可以从 https://github.com/CVPR2023-3D-Occupancy-Prediction/CVPR2023-3D-Occupancy-Prediction 下载数据集。 这里建议将下载好的数据集，解压到nuscense数据集文件夹中的 occ3d/gts 下。 同时，为了提升训练的速度，我们对原始的jpg格式的数据集做了一个打包，将其转换成lmdb格式的数据集。

    python3 tools/dataset_converters/create_data_flashocc.py --src-data-path ./tmp_data/nuscenes/ -o ./tmp_data/nuscenes/occ3d

    上面这条命令会在 ./tmp_data/nuscenes/occ3d 下面生成 nuscenes_infos_train.pkl 和 nuscenes_infos_val.pkl 两个文件, 接下来执行下面命令进行打包：

    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 --need-occ
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 --need-occ

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

    tmp_data
    |-- nuscenes
        |-- metas
        |-- v1.0-trainval 
        |-- occ3d
            | -- gts
            | -- nuscenes_infos_train.pkl
            | -- nuscenes_infos_val.pkl
    |-- occ3d_nuscenes
        |-- train_lmdb
        |-- val_lmdb

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

    模型训练

    数据集准备好之后，就可以开始训练浮点型的Occupancy预测模型了。

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

    python3 tools/train.py --stage "float" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py
python3 tools/train.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

    以上命令分别完成浮点模型和定点模型的训练，其中定点模型的训练需要以训练好的浮点模型为基础，具体内容请阅读 量化感知训练（QAT） 章节的内容。

    导出定点模型

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

    python3 tools/export_hbir.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

    模型验证

    在完成训练之后，可以得到训练完成的浮点、量化或定点模型。和训练方法类似，我们可以用相同方法来对训好的模型做指标验证，得到为 Float 、Calibration 和 Quantized 的指标，分别为浮点、量化和完全定点的指标。

    python3 tools/predict.py --stage "float" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

python3 tools/predict.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

    和训练模型时类似，--stage 后面的参数为 "float"、"calibration" 时，分别可以完成对训练好的浮点模型、量化模型的验证。

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

    python3 tools/predict.py --stage "int_infer" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

    模型推理和结果可视化

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

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

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

    定点模型检查和编译

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

    python3 tools/model_checker.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

    在模型训练完成后，可以通过 compile_perf_hbir 脚本将量化模型编译成可以上板运行的 hbm 文件，同时该工具也能预估在 BPU 上的运行性能：

    python3 tools/compile_perf_hbir.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

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

    训练细节

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

    模型构建

    model = dict(
    type="ViewFusion",
    bev_feat_index=-1,
    bev_upscale=2,
    backbone=dict(
        type="HENet",
        in_channels=3,
        block_nums=[4, 3, 8, 6],
        embed_dims=[64, 128, 192, 384],
        attention_block_num=[0, 0, 0, 0],
        mlp_ratios=[2, 2, 2, 3],
        mlp_ratio_attn=2,
        act_layer=["nn.GELU", "nn.GELU", "nn.GELU", "nn.GELU"],
        use_layer_scale=[True, True, True, True],
        layer_scale_init_value=1e-5,
        num_classes=1000,
        include_top=False,
        extra_act=[False, False, False, False],
        final_expand_channel=0,
        feature_mix_channel=1024,
        block_cls=["GroupDWCB", "GroupDWCB", "AltDWCB", "DWCB"],
        down_cls=["S2DDown", "S2DDown", "S2DDown", "None"],
        patch_embed="origin",
    ),
    neck=dict(
        type="FPN",
        in_strides=[2, 4, 8, 16, 32],
        in_channels=[64, 64, 128, 192, 384],
        out_strides=[16, 32],
        out_channels=[256, 256],
        bn_kwargs=dict(eps=1e-5, momentum=0.1),
    ),
    view_transformer=dict(
        type="LSSTransformer",
        in_channels=256,
        feat_channels=64,
        z_range=(-1.0, 5.4),
        depth=depth,
        num_points=num_points,
        bev_size=bev_size,
        grid_size=grid_size,
        num_views=6,
        grid_quant_scale=grid_quant_scale,
        depth_grid_quant_scale=depth_quant_scale,
    ),
    bev_encoder=dict(
        type="BevEncoder",
        backbone=dict(
            type="HENet",
            in_channels=64,
            block_nums=[4, 3, 8, 6],
            embed_dims=[64, 128, 192, 384],
            attention_block_num=[0, 0, 0, 0],
            mlp_ratios=[2, 2, 2, 3],
            mlp_ratio_attn=2,
            act_layer=["nn.GELU", "nn.GELU", "nn.GELU", "nn.GELU"],
            use_layer_scale=[True, True, True, True],
            layer_scale_init_value=1e-5,
            num_classes=1000,
            include_top=False,
            extra_act=[False, False, False, False],
            final_expand_channel=0,
            feature_mix_channel=1024,
            block_cls=["GroupDWCB", "GroupDWCB", "AltDWCB", "DWCB"],
            down_cls=["S2DDown", "S2DDown", "S2DDown", "None"],
            patch_embed="origin",
            quant_input=False,
        ),
        neck=dict(
            type="BiFPN",
            in_strides=[2, 4, 8, 16, 32],
            out_strides=[2, 4, 8, 16, 32],
            stride2channels=dict({2: 64, 4: 64, 8: 128, 16: 192, 32: 384}),
            out_channels=48,
            num_outs=5,
            stack=3,
            start_level=0,
            end_level=-1,
            fpn_name="bifpn_sum",
            upsample_type="function",
            use_fx=True,
        ),
    ),
    bev_decoders=[
        dict(
            type="FlashOccDetDecoder",
            use_mask=True,
            num_classes=num_classes,
            occ_head=dict(
                type="BEVOCCHead2D",
                in_dim=48,
                out_dim=128,
                Dz=16,
                num_classes=num_classes,
                use_predicter=True,
                use_upsample=True,
            ),
            loss_occ=dict(
                type="CrossEntropyLoss",
                use_sigmoid=False,
                ignore_index=255,
                loss_weight=1.0,
            ),
        ),
    ],
)

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

    数据增强

    跟 model 的定义一样，数据增强的流程是通过在config配置文件中定义 data_loader 和 val_data_loader 这两个dict来实现的，分别对应着训练集和验证集的处理流程, 详细信息见 configs/occ/flashocc_henet_lss_occ3d_nuscenes.py。

    训练策略

    为了训练一个精度高的模型，好的训练策略是必不可少的。对于每一个训练任务而言，相应的训练策略同样都定义在其中的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/henet_tinym_imagenet/float-checkpoint-best.pth.tar",
                allow_miss=True,
                ignore_extra=True,
                ignore_tensor_shape=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        lr=1e-4,
        weight_decay=1e-2,
    ),
    batch_processor=batch_processor,
    num_epochs=num_epochs,
    device=None,
    callbacks=[
        stat_callback,
        grad_callback,
        dict(
            type="CosineAnnealingLrUpdater",
            warmup_len=500,
            warmup_by="step",
            warmup_lr_ratio=1.0 / 3,
            step_log_interval=500,
            stop_lr=2e-4 * 1e-3,
        ),
        metric_updater,
        val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="MeanIOU",
        seg_class=occ3d_seg_class,
        ignore_index=17,
    ),
)

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

    注解

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

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

    量化模型训练

    当我们有了纯浮点模型之后，就可以开始训练相应的定点模型了。和浮点训练的方式一样，我们只需要通过运行下面的脚本就可以得到伪量化模型了，该模型仅使用calibration即可达到目标：

    python3 tools/train.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

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

    calibration_trainer = dict(
    type="Calibrator",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
            ),
            dict(type="Float2Calibration", convert_mode=convert_mode),
            dict(
                type="FixWeightQScale",
            ),
        ],
    ),
    data_loader=calibration_data_loader,
    batch_processor=calibration_batch_processor,
    num_steps=calibration_step,
    device=None,
    callbacks=[
        stat_callback,
        calibration_val_callback,
        calibration_ckpt_callback,
    ],
    val_metrics=dict(
        type="MeanIOU",
        seg_class=occ3d_seg_class,
        ignore_index=17,
    ),
    log_interval=calibration_step / 10,
)

    quantize参数的值不同

    当我们训练量化模型的时候，需要设置quantize=True，此时相应的浮点模型会被转换成量化模型，相关代码如下：

    model.fuse_model()
model.set_qconfig()
horizon.quantization.prepare_qat(model, inplace=True)

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

    训练策略不同

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

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