10.1.4.4. MobileNetV1分类模型训练

这篇教程主要是告诉大家如何利用HAT在 ImageNet 上训练一个state-of-art的浮点和定点模型。 ImageNet 是图像分类里用的最多的数据集,很多最先进的图像分类研究都会优先基于这个数据集做好验证。 虽然有很多方法在社区或者其他途径里可以获取到state-of-art的分类模型,但从头训一个state-of-art的分类模型仍然不是一个简单的任务。 本篇教程将会重点讲叙从数据集准备开始如何在 ImageNet 上训练出一个state-of-art的模型,包括浮点、量化和定点三种模式。

其中 ImageNet 数据集可以从 ImageNet官网 进行下载,下载之后的数据集格式为:

ILSVRC2012_img_train.tar
ILSVRC2012_img_val.tar
ILSVRC2012_devkit_t12.tar.gz

这里我们使用 MobileNetV1 的例子来详细介绍整个分类的流程。

10.1.4.4.1. 训练流程

如果你只是想简单的使用 HAT 的接口来进行简单的实验,那么首先阅读一下这一小节的内容是个不错的选择。 对于所有的训练和评测的任务, HAT 统一采用 tools + config 的形式来完成。 在准备好原始数据集之后,通过以下的流程,我们可以方便的完成整个训练的流程。

10.1.4.4.1.1. 数据集准备

首先是数据集打包,打包数据集与原始数据集在处理速度上有明显的优势,这里我们选择与 PyTorch 一脉相承的 LMDB 的打包方法, 当然由于HAT在处理 dataset 上的灵活性,其他形式的数据集打包和读取形式,如 MXRecord 也是可以独立支持的。

tools/datasets 目录下提供了 cityscapesimagenetvocmscoco 这些常见数据集的打包脚本。 例如 imagenet_packer 的脚本,可以利用 torchvision 提供的默认公开数据集处理方法直接将原始的公开 ImageNet 数据集转成 Numpy 或者 Tensor 的格式, 最后将得到的数据统一用 msgpack 的方法压缩到 LMDB 的文件中。

这个过程可以很方便通过下面的脚本完成数据集打包:

python3 tools/datasets/imagenet_packer.py --src-data-dir ${src-data-dir} --target-data-dir ${target-data-dir} --split-name train --num-workers 10 --pack-type lmdb

python3 tools/datasets/imagenet_packer.py --src-data-dir ${src-data-dir} --target-data-dir ${target-data-dir} --split-name val --num-workers 10 --pack-type lmdb

在完成数据集打包之后,可以得到含有 ImageNetLMDB 数据集。下一步就可以开始训练。

10.1.4.4.1.2. 模型训练

准备打包好数据集之后,便可以开始训练模型。只需要运行下面的命令就可以启动训练:

python3 tools/train.py --stage "float" --config configs/classification/mobilenetv1_imagenet.py

由于HAT算法包使用了注册机制,使得每一个训练任务都可以按照这种 train.py 加上 config 配置文件的形式启动。 train.py 是统一的训练脚本,与任务无关,我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的 config 配置文件里面。 上面的命令中 --stage 后面的参数可以是 "float""calibration""int_infer" ,分别可以完成浮点模型、量化模型的训练以及量化模型到定点模型的转化, 其中量化模型的训练依赖于上一步浮点训练产出的浮点模型,定点模型的转化依赖于量化训练产生的量化模型,具体内容见上一节量化的细节。

10.1.4.4.1.3. 模型验证

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

python3 tools/predict.py --stage "float" --ckpt ${float-checkpoint-path} --config configs/classification/mobilenetv1_imagenet.py

python3 tools/predict.py --stage "calibration" --ckpt ${calibration-checkpoint-path} --config configs/classification/mobilenetv1_imagenet.py

python3 tools/predict.py --stage "int_infer" --ckpt ${int-infer-checkpoint-path} --config configs/classification/mobilenetv1_imagenet.py

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

10.1.4.4.1.4. 模型推理

HAT提供了 infer.py 脚本对各阶段训练好的模型的推理结果进行可视化展示:

python3 tools/infer.py --config configs/classification/mobilenetv1_imagenet.py --model-inputs imgs:${img-path} --save-path ${save_path}

10.1.4.4.1.5. 仿真上板精度验证

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

python3 tools/align_bpu_validation.py --config configs/classification/mobilenetv1_imagenet.py

10.1.4.4.1.6. 定点模型编译

在 HAT 中集成的量化训练工具链主要是为了地平线的计算平台准备的,因此,对于量化模型的检查和编译是必须的。 我们在训练脚本中提供可以对模型检查的接口,可以让用户定义好量化模型之后,先检查能否在 BPU 上正常运行, 目前该检查在启动 train.py 脚本训练 quantize=True 的模型时会默认打开。

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

python3 tools/compile_perf.py --config configs/classification/mobilenetv1_imagenet.py

针对地平线不同架构的BPU,可以在 configs/classification/mobilenetv1_imagenet.py 中设置 march = March.BAYESmarch = March.BERNOULLI2

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

10.1.4.4.1.7. ONNX模型导出

如果想要导出onnx模型, 运行下面的命令即可:

python3 tools/export_onnx.py --config configs/classification/mobilenetv1_imagenet.py

10.1.4.4.2. 训练细节

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

10.1.4.4.2.1. 训练条件

在超过100w张图片的 ImageNet 上训练一个深度学习模型是非常消耗资源的,主要的瓶颈在于矩阵计算和数据读取。

对于矩阵计算,这里非常建议使用一个高性能的GPU来替代CPU作为训练工具,另外同时使用多个GPU可以有效的降低训练时间。

对于数据读取,推荐使用更好的CPU和SSD存储。多线程的CPU加速和更好的SSD存储对于数据读取有很大帮助。 需要注意的是,整个 ImageNet 大概会占用300G的存储资源,所以SSD存储至少得有300G的存储空间。

10.1.4.4.2.2. 网络结构

HAT 或者其他一些社区都可以找到 MobileNetV1 丰富的实现方法,因此, MobileNetV1 的具体实现在这里并不赘述。

在 HAT 的 config 中,我们可以直接用下面的 dict 就可以构建一个浮点的 MobileNetV1 分类模型。 用户可以直接修改 backbone 中的配置参数来达到修改模型目的。

 model = dict(
     type="Classifier",
     backbone=dict(
         type="MobileNetV1",
         num_classes=1000,
         bn_kwargs={},
     ),
     losses=dict(type="CEWithLabelSmooth")
)

模型除了 backbone 之外,还有losses模块,在常见的分类模型中,我们一般直接采用 Cross-Entropy 作为训练的loss, 但是越来越多的实验证明,在分类的loss中加上 Label-Smooth 对训练的结果有帮助,尤其是在结合 Cosine 的lr更新方法上。

通常在定义好一个模型之后,尤其是一些公版模型,我们会有检查计算量的需求。HAT通过 calops 的工具来计算模型的计算量,执行方法如下:

python3 tools/calops.py --config configs/classification/mobilenetv1_imagenet.py

这种计算量的统计工具,是可以同时支持浮点和定点模型的。

10.1.4.4.2.3. 数据增强

关于ImageNet训练的数据增强已经逐渐形成一种共识,我们以 torchvision 提供的数据增强作为基础,构建分类训练的数据增强, 包括 RandomResizedCropRandomHorizontalFlipColorJitter

因为最终跑在BPU上的模型使用的是 YUV444 的图像输入,而一般的训练图像输入都采用 RGB 的形式, 所以HAT提供 BgrToYuv444 的数据增强来将 RGB 转到 YUV444 的格式。 为了优化训练过程,HAT 使用了 batch_processor,可将一些增强处理放在 batch_processor 中优化训练:

dataset=dict(
    type="ImageNetFromLMDB",
    data_path="./tmp_data/imagenet/train_lmdb/",
    transforms=[
        dict(
            type=torchvision.transforms.RandomResizedCrop,
            size=224,
            scale=(0.08, 1.0),
            ratio=(3. / 4., 4. / 3.),
        ),
    ],
    num_samples=1281167
)

对应的 batch_processor 部分:

batch_processor = dict(
    type='BasicBatchProcessor',
    need_grad_update=True,
    batch_transforms=[
        dict(type=torchvision.transforms.RandomHorizontalFlip),
        dict(
            type=torchvision.transforms.ColorJitter,
            brightness=0.4,
            contrast=0.4,
            saturation=0.4,
            hue=0.1,
        ),
        dict(type="BgrToYuv444", rgb_input=True),
        dict(type=torchvision.transforms.Normalize, mean=128.0, std=128.0),
    ],
)

验证集的数据转换相对简单很多,最主要的区别是做短边 Resize 到256,和 CenterCrop 。其他的颜色空间转换和训练集是一样的。

dataset=dict(
    type="ImageNetFromLMDB",
    data_path="./tmp_data/imagenet/val_lmdb/",
    transforms=[
        dict(type=torchvision.transforms.Resize, size=256),
        dict(type=torchvision.transforms.CenterCrop, size=224),
    ],
    num_samples=50000,
),

val_batch_processor = dict(
    type='BasicBatchProcessor',
    need_grad_update=False,
    batch_transforms=[
        dict(type="BgrToYuv444", rgb_input=True),
        dict(type=torchvision.transforms.Normalize, mean=128.0, std=128.0),
    ],
)

10.1.4.4.2.4. 训练策略

ImageNet 上训练不同分类模型的训练策略大体上一致,但也有微小的区别。这里我们主要介绍有效果提升的细节。

Cosine 的学习策略配合 Warmup 与普通的 StepLr 有一些提升效果。适当延长epoch的训练长度对于小模型也有提升。

另外,只对 weight 的参数施加 L2 norm 也是推荐的训练策略。 configs/classification/mobilenetv1_imagenet.py 文件中的 float_trainercalibration_trainerint_trainer 分别对应浮点、量化、定点模型的训练策略。下面以 float_trainer 训练策略为例:

float_trainer = dict(
     type='distributed_data_parallel_trainer',
     model=model,
     data_loader=data_loader,
     optimizer=dict(
         type=torch.optim.SGD,
         params={"weight": dict(weight_decay=4e-5)},
         lr=0.4,
         momentum=0.9,
     ),
     batch_processor=batch_processor,
     num_epochs=240,
     device=None,
     callbacks=[
         stat_callback,
         dict(
             type="CosLrUpdater",
             warmup_by='epoch',
             warmup_len=5,
         ),
         metric_updater,
         val_callback,
         ckpt_callback,
     ],
     train_metrics=[ dict(type="Accuracy"), dict(type="TopKAccuracy", top_k=5)],
     val_metrics=[ dict(type="Accuracy"), dict(type="TopKAccuracy", top_k=5)],
)

10.1.4.4.2.5. 量化训练

关于量化训练中的关键步骤,比如准备浮点模型、算子替换、插入量化和反量化节点、设置量化参数以及算子的融合等, 请阅读 量化感知训练 章节的内容。 这里主要讲一下HAT的分类中如何定义和使用量化模型。

在模型准备的好情况下,包括量化已有的一些模块完成之后,HAT在训练脚本中统一使用下面的脚本将浮点模型映射到定点模型上来。

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

量化训练的策略并不统一,这里简单描述分类模型训练中的常见策略。

量化训练的整体策略可以直接沿用浮点训练的策略,但学习率和训练长度需要适当调整。因为有浮点预训练模型,所以量化训练的学习率 Lr 可以很小, 一般可以从0.001或0.0001开始,并可以搭配 StepLrUpdater 做1-2次 scale=0.1Lr 调整;同时训练的长度不用很长。 此外 weight decay 也会对训练结果有一定影响。

10.1.4.4.2.6. 模型检查编译和仿真上板精度验证

对于HAT来说,量化模型的意义在于可以在 BPU 上直接运行。因此,对于量化模型的检查和编译是必须的。 前文提到的 compile_perf 脚本也可以让用户定义好量化模型之后,先检查能否在 BPU 上正常运行, 并可通过 align_bpu_validation 脚本获取模型上板精度。用法同前文。

10.1.4.4.3. 预训练模型

HAT 已经提供了丰富的在 ImageNet 上的预训练模型,可以参照 modelzoo 的内容,所有模型都在发布包中。