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/occ3d_nuscenes_packer.py --src-data-dir ./tmp_data/nuscenes/ --pack-type lmdb --target-data-dir ./tmp_data/occ3d_nuscenes --split-name test --anno-file ./tmp_data/nuscenes/occ3d/nuscenes_infos_val.pkl python3 tools/datasets/occ3d_nuscenes_packer.py --src-data-dir ./tmp_data/nuscenes/ --pack-type lmdb --target-data-dir ./tmp_data/occ3d_nuscenes --split-name train --anno-file ./tmp_data/nuscenes/occ3d/nuscenes_infos_train.pkl

上面这两条命令分别对应着转换训练数据集和验证数据集,打包完成之后,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 python3 tools/train.py --stage "qat" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

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

导出定点模型

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

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

模型验证

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

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 imagedir:${img-path},homo:${homo-path} --save-path ${save_path}

仿真上板精度验证

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

python3 tools/validation_hbir.py --stage "align_bpu" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

定点模型检查和编译

在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_loaderval_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)

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

训练策略不同

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

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