【我参加NVIDIA Sky Hackathon】使用TAO模型训练工具训练CV检测识别模型

目录

1.TAO模型训练工具

2.环境配置

3.CV模型框架选择

4.数据集处理

4.1.数据收集

4.2.数据标注

5.模型训练

5.1.SSD

5.2.yolov4

5.3.模型剪枝

6.模型部署

6.1.模型导出

6.2.模型转换

---

1.TAO模型训练工具

        TAO是由NVIDIA提供的一款开元、便捷的模型训练工具，主要用于做迁移学习。其收录了许许多多预训练模型，可以帮助开发者快速打造出自己需要的模型。

通过NVIDIA举办的第七届Sky Hackathon比赛，我体验到了使用TAO模型训练工具训练CV检测识别模型整个流程。

2.环境配置

        想要使用TAO模型训练工具训练视觉模型，需要安装大量前置库。这些前置库的列表与安装方式网上已有详细说明，在此不做说明。

        本次比赛我们使用的系统为Ubuntu-18.04系统，python版本为3.8，同时配置了CUDA-11.4，docker与nvidia-docker2，MiniConda3与JupyterLab。

3.CV模型框架选择

        TAO模型训练工具收录了大量预训练模型，且这些模型对于各种各样的框架都基本兼容。本次比赛我们尝试使用TAO训练了ssd-resnet18和yolov4-resnet18。

4.数据集处理

4.1.数据收集

        本次比赛的主题为垃圾分类，主要任务为将图片中的垃圾检测出来，并将类分类为果皮、纸箱和瓶子。所以本次我们需要收集大量有关垃圾的图片数据。

        本次比赛官方已经提供了2600多张图片数据，为我们的数据收集工作减轻了些许负担。在对官方提供的数据做筛选后，还剩下2000来张良好的图片数据。

        除了官方提供的数据集，我们还通过网络爬虫，实地采集获得了400多张图片数据。

        整合完所有数据后，我们认为该数据集的数据量依然不够大，于是我们利用opencv对所有数据做了数据增强，增强方式为对所有数据进行随机的平移、旋转和翻转，以及以十分之一的概率对原数据集进行随机地加入高斯噪声。最终，我们获得了一个数据量为4800左右的数据集。

4.2.数据标注

        使用TAO模型训练工具进行CV检测识别模型的训练需要的数据标注类型为kitti数据集，这也是使用TAO训练模型的便利之处。同样的kitti数据集，我们可以将其用在TAO收录的所有不同的模型框架上，不需要再重新对数据进行标注。

        关于对数据的标注，我们使用的标注工具是labelimg。通过分工，我们花了大概三天时间对4800多张图片进行了标注。

5.模型训练

        使用TAO训练模型首先需要NGC密钥，该密钥可从NVIDIA官网获取。密钥的作用是很好地保护了开发者的知识产权，脱离了密钥，模型无法被转换，也无法使用推理。

        使用TAO训练模型需要通过写specs文件来对训练参数，各种路径进行定义，这些文件的模板在官网均有详细说明，网址为TAO Toolkit — TAO Toolkit 3.22.05 documentation (nvidia.com)。

       使用TAO训练模型还需要下载预训练模型，所有的预训练模型都收录在官网，网址为TAO Pretrained Object Detection | NVIDIA NGC。

5.1.SSD

        ssd-resnet18模型框架为比赛官方推荐并提供了使用教程的框架。其训练成本较低，梯度下降速度较快，使用8g显存的rtx3070显卡对其进行训练，batch-size可以设置到32。但该框架容易过拟合，经过80轮训练后val-loss依旧有30多。且fps远远不如yolo框架。

        在导入好环境，定义好参数后，使用tao ssd train命令来进行训练。

5.2.yolov4

        Yolov4-resnet18模型框架为比赛加分框架。官方没有提供太多的教程，且其训练成本比ssd-resnet18要高，使用8g显存的rtx3070显卡对其进行训练，batch-size只能设置为4。且其梯度下降速度较慢，训练时长较久。但其精度较高，经过80轮训练后val-loss大概在5左右，且其fps要比ssd模型要高。

        在导入好环境，定义好参数后，使用tao yolo_v4 train命令来进行训练。

        在训练yolov4模型的过程中，我们发现了一个关于训练进度显示的一个小bug。具体为当其specs文件中的include_difficult_in_training变量设为True时，每次更新进度都将会显示如下图的提示。

 

        该提示末尾并没有设置不换行输出，如果开发者用jupyter lab进行训练，且训练批次较大时，浏览器页面将会因为内存溢出而页面崩溃。如图。

        解决方法为将该变量设为True；

        或是找到iostream.py文件，将self.max_buffer_size设大点。

5.3.模型剪枝

       模型剪枝操作是指对模型各层的神经元进行随机的摘除，这样做可以减少各神经元之间的依赖性，能更好地提取特征值，从而防止过拟合。同时也能使模型更加精简，增大模型的fps值。

       但剪枝操作也会降低模型的精度，建议在确保模型精度较高时再进行剪枝操作。

       由于我们本次训练出来的模型精度不算特别高，所以为确保精度，我们决定不进行剪枝操作。

6.模型部署

6.1.模型导出

    模型训练好后，将会以tlt的格式保存，我们可以通过观察训练日志，挑选出mAP值最高的模型，并利用tao ssd export（yolov4模型使用tao yolo_v4 export）指令将该tlt格式的模型文件转换为etle格式的文件并导出。

6.2.模型转换

       导出模型后，我们可以将模型转移到我们想将其部署到的平台上。在平台上搭建好环境（环境配置在官方文档以及各论坛上有详细说明）后，需要在官方网站下载对应该平台环境的tao-converter文件（网址为TAO Converter | NVIDIA NGC），便可使用命令tao converter将etlt格式的模型转换为engine格式或bin格式的推理器。（tao converter命令的使用方式在官方文档有详细说明）

    紧接着，我们要将需要推理的图片数据进行预处理再送入推理器推理，经过推理得到的数据还需要经过后处理来表达。预处理和后处理的脚本可以在网上搜寻到，也可以自己编写。

    本次比赛由于我们暂时无法解决yolov4的后处理问题，考虑到时间成本，只能决定放弃使用yolov4-resnet18模型，而去使用ssd-resnet18模型。