4.7.5.6. Unet分割模型训练

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

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

1. 训练流程

数据集下载

要下载 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 提供了cityscapes2lmdb.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.py 配置文件中的数据集路径（data_rootdir）设置为数据集 LMDB 文件所在位置。

在网络训练开始之前，你可以使用以下命令先测试一下网络的计算量和参数数量：

python3 tools/calops.py --config configs/segmentation/unet.py --input-shape "1,3,1024,2048"

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

python3 tools/train.py --step float --config configs/segmentation/unet.py

如果想要验证已经训练好的模型精度,运行下面的命令即可：

python3 tools/train.py --step float --config configs/segmentation/unet.py --val-ckpt float-checkpoint-best.pth.tar --val-only

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

python3 tools/export_onnx.py --config configs/segmentation/unet.py --ckpt float-checkpoint-best.pth.tar --onnx-name unet.onnx