#UNet分割模型训练

作为HAT分割任务的示例，这篇教程主要向大家展示如何使用HAT在 Cityscapes 数据集上训练一个 state-of-the-art 的浮点和定点模型。

Cityscapes 是一个城市驾驶场景的图像数据集，包含了5000张具有像素级标注的图像，图像中的物体被分为19个类别。 分割任务相对复杂，对模型能力的要求较高，使用小模型在分割任务上取得较好的指标并不是很轻易的事。 本教程将从零开始，详细描述如何使用HAT在 Cityscapes 数据集上训练一个 state-of-the-art 的分割模型，并在浮点模型基础上进行量化训练，最终得到一个定点模型。

#训练流程

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

#数据集准备

#数据集下载

要下载 Cityscapes 数据集，首先需要在 官方网站 注册一个账号。

之后便可在 下载页面 下载需要的数据集文件， 这里我们只需要 gtFine_trainvaltest.zip 和 leftImg8bit_trainvaltest.zip 两个文件。

同时， Cityscapes 数据集官方还提供了一个脚本用于数据的下载和处理，见 Github链接 。 首先使用如下命令安装官方工具：

python3 -m pip install cityscapesscripts

然后使用官方工具下载所需数据集文件（注意，使用此工具下载仍需要登录上面注册的账号）。

csDownload -d ${data-dir} gtFine_trainvaltest.zip
csDownload -d ${data-dir} leftImg8bit_trainvaltest.zip

最后对下载好的文件进行解包即可（可选）：

cd ${data-dir}
unzip gtFine_trainvaltest.zip
unzip leftImg8bit_trainvaltest.zip

#数据集打包

为了高效地读取数据，我们推荐预先将数据集打包为 LMDB 格式。 HAT提供了 cityscapes_packer.py 脚本来方便地将数据集从原始公开的格式转换为 numpy.ndarray或 torch.Tensor ，使用 msgpack 对数据进行封装，并最终打包为 LMDB 文件。

数据集打包的命令如下：

python3 tools/datasets/cityscapes_packer.py --src-data-dir ${data-dir} --split-name train --pack-type lmdb
python3 tools/datasets/cityscapes_packer.py --src-data-dir ${data-dir} --split-name val --pack-type lmdb

生成的lmdb文件保存在 ${data-dir}/train_lmdb 和 ${data-dir}/val_lmdb 路径下。

#模型训练

将数据集打包为 LMDB 文件后，就可以开始模型的训练了。HAT提供了 train.py训练脚本来方便地配合 config 文件实现模型训练。

模型训练的命令如下，开始训练之前请确保将 unet_mobilenetv1_cityscapes.py 配置文件中的数据集路径（data_rootdir）设置为数据集 LMDB 文件所在位置。

python3 tools/train.py --stage "float" --config configs/segmentation/unet_mobilenetv1_cityscapes.py
python3 tools/train.py --stage "calibration" --config configs/segmentation/unet_mobilenetv1_cityscapes.py
python3 tools/train.py --stage "qat" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

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

#导出定点模型

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

python3 tools/export_hbir.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

#模型验证

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

python3 tools/predict.py --stage "float" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

python3 tools/predict.py --stage "calibration" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

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

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

python3 tools/predict.py --stage "int_infer" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

#模型推理和结果可视化

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

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

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

#定点模型检查和编译

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

python3 tools/model_checker.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

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

python3 tools/compile_perf_hbir.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

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

#训练细节

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

#模型构建

model = dict(
    type="Segmentor",
    backbone=dict(
        type="MobileNetV1",
        num_classes=-1,
        bn_kwargs=bn_kwargs,
        alpha=alpha,
        dw_with_relu=True,
        include_top=False,
        flat_output=False,
    ),
    neck=dict(
        type="DwUnet",
        base_channels=int(32 * alpha),
        bn_kwargs=bn_kwargs,
        act_type=nn.ReLU,
        use_deconv=False,
        dw_with_act=True,
        output_scales=train_scales,
    ),
    head=dict(
        type="SegHead",
        num_classes=num_classes,
        in_strides=train_scales,
        out_strides=train_scales,
        stride2channels={
            stride: int(stride * 32 * alpha) for stride in train_scales
        },
        feat_channels=tuple(np.array(train_scales) * int(32 * alpha)),
        stacked_convs=0,
        argmax_output=False,
        dequant_output=True,
        int8_output=True,
        upscale=False,
        output_with_bn=output_with_bn,
        bn_kwargs=bn_kwargs,
    ),
    losses=dict(
        type="SoftmaxFocalLoss",
        loss_name="Focal",
        num_classes=num_classes,
        weight=tuple(np.array((256, 128, 64, 32, 16)) / 19),
        reduction="mean",
    ),
)

分割模型主要由 backbone， neck ， head 三部分组成，这里我们使用了 MobileNetV1_0.25 作为backbone，MobileNet是一种轻量高效的网络结构。neck则使用了 Unet 结构，这种结构可以综合各个尺度上的 featuremap，保留精细的空间信息。 head 为一个卷积层，负责输出最终的分割结果。

我们使用了 FocalLoss 作为损失函数。 FocalLoss 可视为一种动态权重的交叉熵损失函数，可以较好地解决类别不平衡带来的训练困难的问题。

Unet 的层次化结构和 FPN 的思想一致，很适合采用相同的训练方式，即在 Unet 的每个 scale 上都构造一个输出， 将此输出和对应尺寸的 ground truth 构造损失函数进行训练，通过对网络的各个尺度都进行监督，为网络训练提供更加丰富的参考信息，降低训练难度，提高训练速度和最终精度。 同时考虑到我们需要的最终结果是最大尺寸网络输出(scale=4)，为了避免其他尺寸的梯度过大影响到最大尺寸输出的精度，我们按照 scale 为损失函数增加了相应的权重，scale 越大的层权重越小。

通常在定义好一个模型之后，尤其是一些公版模型，我们会有检查计算量的需求。HAT提供了 calops.py 来计算模型的计算量，命令如下：

python3 tools/calops.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

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

#数据预处理

首先，我们使用 LabelRemap 将数据标签重映射至 [0, 18] 的区间。

对于训练集， SegRandomAffine 可以对图片进行随机仿射变换来进行数据增强，我们只配置了随机的缩放，没有做任何旋转操作。

由于训练采用了类 FPN 的方式，我们需要将标签缩放至不同的大小，用于模型不同 scale 的训练。

因为最终跑在BPU上的模型使用的是 YUV444 的图像输入，而一般的训练图像输入都采用 RGB 的格式， 所以HAT提供 ImageBgrToYuv444 的数据增强来将 RGB 数据转到 YUV444 的格式。

最后，归一化对于深度学习模型训练是必须的。

值得注意的是，这里我们使用了 MultiBatchProcessor 进行任务的训练，此 Processor 支持在GPU上以 batch 为单位对数据做预处理。 由于分割任务的数据预处理过程相对复杂，若使用CPU来做，会出现处理瓶颈。使用GPU预处理数据会提高显存使用量，导致最大 batch size 降低。 但即使如此，最终的训练速度也有可观的提升。

data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="Cityscapes",
        data_path=os.path.join(data_rootdir, "train_lmdb"),
        transforms=[
            dict(type="PILToTensor"),
        ],
    ),
    sampler=dict(type=torch.utils.data.DistributedSampler),
    batch_size=batch_size_per_gpu,
    shuffle=True,
    num_workers=dataloader_workers,
    pin_memory=True,
    collate_fn=collate_2d,
)
batch_processor = dict(
    type="BasicBatchProcessor",
    need_grad_update=True,
    batch_transforms=[
        dict(type="LabelRemap", mapping=CITYSCAPES_LABLE_MAPPINGS),
        dict(type="SegOneHot", num_classes=num_classes),
        dict(type="SegResize", size=data_shape[1:]),
        dict(
            type="SegRandomAffine",
            degrees=0,
            scale=(0.5, 2.0),
            interpolation=InterpolationMode.BILINEAR,
            label_fill_value=0,
        ),
        dict(type="BgrToYuv444", rgb_input=True),
        dict(
            type="TorchVisionAdapter",
            interface="Normalize",
            mean=128.0,
            std=128.0,
        ),
        dict(
            type="Scale",
            scales=tuple(1 / np.array(train_scales)),
            mode="bilinear",
        ),
    ],
    loss_collector=loss_collector,
    enable_amp=enable_amp,
)

验证集的数据预处理与训练集相比不需要做随机的仿射变换和多尺度缩放，其他步骤一致，不再赘述。

#训练策略

分割任务可以看作是像素级的分类任务，因此训练策略也和分类任务高度相似，首先在保证可以收敛的前提下尽可能调大学习率可以提高训练速度； 在使用某一学习率训练至精度不再增长时，适当将学习率调小，模型便可继续收敛，精度进一步提升； 最终训练完成后对测试集精度和训练集精度进行比较，若训练集精度高于测试集过多，则可以认为模型出现了过拟合，此时适当增大 weight decay 可以增强模型的泛化能力，减少过拟合，从而获得更高的测试集精度。

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=weight_decay)},
        lr=start_lr,
        momentum=0.9,
    ),
    batch_processor=batch_processor,
    device=None,
    num_epochs=train_epochs,
    callbacks=[
        stat_callback,
        dict(
            type="StepDecayLrUpdater",
            lr_decay_id=[200, 240, 280],
            lr_decay_factor=0.1,
        ),
        metric_updater,
        tb_callback,
        aidi_tb_callback,
        tb_loss_callback,
        aidi_tb_loss_callback,
        float_val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
)

#量化模型训练

量化训练的目的是在训练好的浮点模型基础上对数据进行模拟量化，来模拟定点计算的过程，使得经过量化训练的模型转为定点模型时的精度损失降到最低。

由于浮点模型经过了充分的训练已经收敛到了较优的状态，量化训练通常只需要对模型进行微调即可，学习率不能设置过大，可以从 1e-4 的数量级开始尝试，其他参数一般和浮点训练一致即可。同样地，微调模型所需的训练量较低，一般训练几十个 epoch 就可以了。

由于量化训练之前模型已经有一个不错的精度，精度提升的区间较小，且对数据进行了模拟量化，会导致训练过程波动较大。 此时需要一些耐心，认真观察，从波动中看出趋势，适量调整参数，才能取得最好的结果。

#预训练模型

HAT已经提供了此例子的预训练模型，所有模型都在发布包中。