这篇教程以FCOS3D-efficientnetb0为例,告诉大家如何使用HAT算法包训练一个定点的3D检测模型。 在开始量化感知训练,也就是定点模型训练之前,首先需要训练一个精度较高的纯浮点模型,然后基于这个纯浮点模型做finetune,就可以快速的训练出定点模型。 所以我们从训练一个纯浮点的FCOS3D-efficientnetb0模型开始讲起。
在开始训练模型之前,第一步是需要准备好数据集。这里我们训练 FCOS3D 模型使用的是开源的 nuscenes 数据集。解压缩之后数据目录结构如下所示:
同时,为了提升训练的速度,我们对原始的jpg格式的数据集做了一个打包,将其转换成lmdb格式的数据集。只需要运行下面的脚本,就可以成功实现转换:
上面这两条命令分别对应着转换训练数据集和验证数据集,打包完成之后,data目录下的文件结构应该如下所示:
train_lmdb和val_lmdb就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集。
数据集准备好之后,就可以开始训练浮点型的FCOS3D-efficientnetb0检测网络了。 如果你只是单纯的想启动这样的训练任务,只需要运行下面的命令就可以:
由于HAT算法包使用了一种巧妙的注册机制,使得每一个训练任务都可以按照这种train.py加上config配置文件的形式启动。 train.py是统一的训练脚本,与任务无关,我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的config配置文件里面。 config文件里面提供了模型构建、数据读取等关键的dict。
FCOS3D的网络结构可以参考 论文, 这里不做详细介绍。我们通过在config配置文件中定义 model 这样的一个dict型变量,就可以方便的实现对模型的定义和修改。
其中, model 下面的 type 表示定义的模型名称,剩余的变量表示模型的其他组成部分。
这样定义模型的好处在于我们可以很方便的替换我们想要的结构。
跟 model 的定义一样,数据增强的流程是通过在config配置文件中定义 data_loader 和 val_data_loader 这两个dict来实现的,分别对应着训练集和验证集的处理流程。以 data_loader 为例:
为了训练一个精度高的模型,好的训练策略是必不可少的。对于每一个训练任务而言,相应的训练策略同样都定义在其中的config文件中,从 float_trainer 这个变量就可以看出来。
float_trainer从大局上定义了我们的训练方式,包括使用多卡分布式训练(distributed_data_parallel_trainer),模型训练的epoch次数,以及优化器的选择。
同时 callbacks 中体现了模型在训练过程中使用到的小策略以及用户想实现的操作,包括学习率的变换方式(WarmupStepLrUpdater),在训练过程中验证模型的指标(Validation),以及保存(Checkpoint)模型的操作。当然,如果你有自己希望模型在训练过程中实现的操作,也可以按照这种dict的方式添加。
通过上面的介绍,你应该对config文件的功能有了一个比较清楚的认识。然后通过前面提到的训练脚本,就可以训练一个高精度的纯浮点的检测模型。 当然训练一个好的检测模型不是我们最终的目的,它只是做为一个pretrain为我们后面训练定点模型服务的。
当我们有了纯浮点模型之后,就可以开始训练相应的定点模型了。和浮点训练的方式一样,我们只需要通过运行下面的脚本就可以训练定点模型了:
可以看到,我们的配置文件没有改变,只改变了 stage 的类型。此时我们使用的训练策略来自于config文件中的 qat_trainer 和 calibration_trainer 。
当我们训练量化模型的时候,需要设置相应的model_convert_pipeline,此时相应的浮点模型会被转换成量化模型,相关代码如下:
关于量化训练中的关键步骤,比如准备浮点模型、算子替换、插入量化和反量化节点、设置量化参数以及算子的融合等,请阅读 量化感知训练 章节的内容。
正如我们之前所说,量化训练其实是在纯浮点训练基础上的finetue。因此量化训练的时候,我们的初始学习率设置为浮点训练的十分之一,训练的epoch次数也大大减少,最重要的是 model 定义的时候,我们的 pretrained 需要设置成已经训练出来的纯浮点模型的地址。
做完这些简单的调整之后,就可以开始训练我们的量化模型了。
完成量化训练后,便可以开始导出定点模型。可以通过下面命令来导出:
模型训练完成之后,我们还可以验证训练出来的模型性能。由于我们提供了float,calibration和qat三阶段的训练过程,相应的我们可以验证这三个阶段训练出来的模型性能,只需要相应的运行以下两条命令即可:
同时,我们还提供了quantization模型的性能测试,但需要注意是必须要先导出hbir:
这个显示出来的精度才是最终的int8模型的真正精度,当然这个精度和qat验证阶段的精度应该是保持十分接近的。
除了上述模型验证之外,我们还提供和上板完全一致的精度验证方法,可以通过下面的方式完成:
如果你希望可以看到训练出来的模型对于单张图片的检测效果,我们的tools文件夹下面同样提供了单张图片预测及可视化的脚本,你只需要运行以下脚本即可:
在训练完成之后,可以使用 compile_perf_hbir 的工具用来将量化模型编译成可以上板运行的 hbm 文件,同时该工具也能预估在BPU上的运行性能,可以采用以下脚本: