Deeplab分割模型训练

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

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

1. 训练流程

数据集下载

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

之后便可在 下载页面 下载需要的数据集文件,这里我们只需要 gtFine_trainvaltest.zipleftImg8bit_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.ndarraytorch.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 文件实现模型训练。

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

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

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

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

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

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

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

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

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