物件检测的模型训练与优化

物件检测 (objectdetection) 或物件定位 (object location) 是目前普及度最高的人工智能应用，也是过去几年中神经网络发展最迅猛的领域，短短几年之中创造出非常多优异的经典算法，并且各有所长，如今在数量与质量上都还持续进化中，这现象其实也是造成使用者很大困扰之处，因为究竟哪种神经网络最适合您的使用场景，最终还是得经过实际验证之后才能确认。

TAO 工具非常高效地协助我们解决这个难题，让我们在不了解各种神经网络结构与算法、也不熟悉 Tensorflow 或 Pytorch 这些复杂框架之前，不用撰写任何 C++ 或 Python 代码，就能非常轻易地训练好深度学习模型，并且快速验证不同网络模型的性能与效果，这估计能减少能降低 80% 以上的前期学习与评估的时间成本。

在 TAO 的 cv_samples 里提供多个物件检测算法，包括 detectnet_v2、dssd、ssd、facenet、fast_rcnn、lprnet、mask_rcnn、retinanet、unet、ssd、yolo_v4等十多个物件检测范例，每一个流程与本系列前一个 “图像分类的模型训练” 的范例是一致的，因此这里只挑选目前普及度较高的 ssd 神经网络进行示范就行。

首先进入 Jupyter 操作界面，打开 cv_samples/ssd/ssd.ipynb脚本，然后按照以下步骤逐步执行：

0. 环境变量配置与路径映射（Setup env variables and map drives）：

在系列文章 “视觉类训练脚本的环境配置” 内有详细的讲解，这里只要对下面两变量进行处理就可以：

• %env KEY=<您的NGC秘钥>

• 将“%envLOCAL_PROJECT_DIR=YOUR_LOCAL_PROJECT_DIR_PATH”改成“os.environ["LOCAL_PROJECT_DIR"] = os.getcwd()”

然后执行这部分的 4 个指令块

1. 安装 TAO 启动器（Install the TAO launcher）：

跳过第一个安装指令块，执行“!tao info”指令块就行。

2. 准备数据集与预训练模型（Preparedataset and pre-trained model）

数据集处理通常是模型训练过程比较消耗人力的步骤。这里推荐不使用脚本里的步骤去下载与解压缩，因为这里指向的下载位置在国外网站，这对大部分使用者是很困扰的，因此最好自行手动下载处理。

(1) TAO 的物件检测模型训练支持 KITTI 标注格式，这里使用的范例数据集与系列文章“视觉数据格式处理与增强”是同一套，如果前面实验已将下载的话就复制一份过来就行，也可以重新在https://pan.baidu.com/s/1WcRt9ONqsYhLRmL8Rg8_YQ（密码855n），下载 data_object_image_2.zip 与 data_object_label_2.zip 到 ssd.ipynb 相同目录下，然后执行以下指令进行解压缩：

cd ~/tao/cv_samples/ssd  # 请根据您实际路径unzip -u data_object_image_2.zip -d dataunzip-udata_object_label_2.zip-ddata

会在这个工作目录下生成 data/training/image_2、data/training/label_2 与 data/testing三个子目录，其中训练用数据量为 7481 张图片、测试数据量为 7518 张。

如果有自己的数据集想要转到 TAO 来进行训练，除了先将标注格式转成 KITTI 之外，最好也遵循上面的路径结构，这样就不需修改 spec 目录下相关配置文件里面的设定值，对初学者是最便利的方式。

(1) 接着调用 generate_val_dataset.py，从训练数据中提取 10% 内容作为校验用途，存放在 data/val 目录，于是在 data/training 下面的图像与标注文件数量就剩下 6733 个。然后执行 “2.3Generate tfrecords” 指令块，将训练与校验数据集都转成 tfrecords 格式，分别存放在对应的路径之下。

(2) 如果要配合迁移学习功能的话，接下来就是从 NGC 下载合适的预训练模型。这里选择的是 “nvidia/tao/pretrained_object_detection:resnet18”，最终下载的模型文件为 89MB 大小的 resnet_18.hdf5 模型文件。

只要这个数据集整理步骤做得完整，后面的部分就会很简单。

3. 提供训练用配置文件（Providetraining specification）

在 TAO 所有范例里，配置文件是所有内容的精华之处，这必须由专业算法工程师群进行调试，不仅要这些神经网络的结构与算法有足够深入的了解，还需要经过足够长的重复调试过程，才能调试出一个够完整的优化参数组合，这部分是英伟达非常庞大精英人力与成本的所提炼出来的精髓。

物件检测应用的设定组，比前面的图像分类要复杂很多，虽然初学者没有能力去做细节参数的调整，但还是得了解一下主要结构，对于学习深度神经网络会有很大的帮助。

在 ssd 项目的 specs 目录下有 5 个配置文件，其中与训练相关的有 4 个，但实际使用到的就是ssd_train_resnet18_kitti.txt与ssd_retrain_resnet18_kitti.txt这两个配置文件。

前者是提供给首次训练模型时使用，后者是对修剪过的模型进行再训练时使用，两者的内容几乎一样，主要差异在配合迁移学习功能的部分，前者使用“pretrain_model_path”指向从NGC下载的预训练模型，后者使用“pruned_model_path”指向后面会操作的 “修剪过的模型”，其他参数都一样。

现在就以 ssd_train_resnet18_kitti.txt 配置文件来做说明，里面总共有 6 个参数组：

(1) ssd_config：

这是根据网络种类所命名，不同网络的配置组内容都不太一样，例如 yolov3 范例的配置组为 “yolov3_config”，以此类推。每种神经网络都有截然不同的参数内容，在 ssd 网络有 12 个参数的设定，而 yolo_v4 网络有将近 20 个，所以网络之间的配置是不能套用或互换的。如果未来要开发自己的特殊模型，最好从现有的 20 范例里面去挑选，就能以该网络的配置文件为基础进行修改。

关于这个配置组的内容就不多时间去说明，初学者先不要贸然去修改，以免训练过程发生不可预期的错误，或者模型效果不好。

(2) training_config：

执行训练时的一些参数设定，这是使用者能进行调整的部分，包括以下部分：

• batch_size_per_gpu：根据执行训练的GPU总显存进行调整

• num_epochs：训练回合数，这里设定为80

• enable_qat：开启QAT训练模式，这里预设值为 ”false”，表示不开启

• checkpoint_interval：每多少回合生成一次模型文件，这里预设值为1

• learning_rate、regularizer 与 optimizer：如果对 ssd 有深入了解的，可以进一步调整。

(3) evaluation_config：

这些存放模型评估的一些参数，可以根据计算设备的显存值调整 batch_size 大小，不过这些并不影响训练的结果。

(4) nms_config：

NMS (Non-maximumsuppression) 非极大抑制应用的目的，是要消除多余的框，找到最佳的物体检查位置。

(5) augmentation_config：

这部分的参数内容请参考系类文章 “视觉类的数据增强”，有详细的说明。虽然没办法进入 TAO 源码（未开放）去确认，但这项参数设定组应该显示在物件检测模型训练过程中，并将“数据增强”功能加进来，这对提高模型精准度会起到很大的效果。

(6) dataset_config：

这里最主要配置数据源的路径，以及类别 “key:value” 对的映射关系；

• 一组 data_source：

里面包含 tfrecords_path 与 image_directory_path 两个路径设定；

• 多组 target_class_mapping：

每一组对应的 “key” 是数据集的原始类别名称，后面映射的 “value” 是训练后模型的识别分类。

例如这里使用的 KITTI 数据集中，原始分类有 “car/ van/cyclist/ person_sitting/ pedestrian/ DontCare/ tram/ truck/ misc” 等 9 种分类，经过映射后只保留 “Car/cyclist/ pedestrian” 三个分类，没有在映射组里的分类就会忽略，例如 truck、misc 等。

如果使用自己的数据集训练模型，这部分需要自行调整到合适的类别对应，否则训练的结果也不能实现您要达到的识别效果。

• 一组 validation_data_sources：

包括 label_directory_path 与 label_directory_path，存放校验数据集的图像与标注路径。

以上 6 组配置在每个物件检测模型的配置文件中都有，但是配置内容不尽相同， TAO 已经为所支持的网络都提供优化过的参数内容，初学者只要直接使用就可以，等日后有更深入的掌握，再试着去修改这些参数。

原文标题：NVIDIA Jetson Nano 2GB 系列文章（61）：物件检测的模型训练与优化

文章出处：【微信公众号：NVIDIA英伟达企业解决方案】欢迎添加关注！文章转载请注明出处。

审核编辑：汤梓红