Jetson Nano 2GB 系列文章（25）： “Hello AI World”图像分类的模型训练

前面的所有内容与实验，全部都是关于“推理识别（inference）”的部分，主要是为了让用户先体会到“最终结果”，包括需要花费的精力、以及最后得到的效率。由于项目提供非常多的预训练模型文件，因此我们可以很轻松地调用这些资源，去进行推理计算，并且得到相当惊人的性能。

然而，真正要完成一个深度学习的项目，还是必须完成前面阶段的模型训练部分，本文的重点就是带着大家，利用Hello AI World所提供的资源，开始进行简单的图像分类应用的模型训练。

所有深度学习的模型训练，都必须具备以下的几个元素，缺一不可：

• 数据集（dataset）:

事实上数据集的质量，直接关系到训练出来的模型质量，就如同做菜一样，如果食材不够好，再厉害的厨师也没办法做出真正的好菜。

数据集的来源主要有三个：

• 网上已有的开源资源：包括ImageNet、MS COCO、OpenImages等等，都是数量级在百万张图片以上的数据集，而且分类十分清楚，是个非常重要的资源，必须学会如何从这些数据集中，萃取自己所需要的的数据。
• 自行收集：通过在网上根据关键字去搜索下载，也可以利用自己的摄像设备去拍摄收集。
• 自行转换：可以将已有的图像数据，经过灰度处理、镜像映射等方式，再生成2x~4x数据量。

• 深度神经网络模型（算法）：

这个部分是深度学习的精髓，目前在几个主流神经网络模型之间，很难去评孰优孰劣，就像大餐厅的名厨一样，各有侧重点与风格。在这里，为了要能在Jetson设备上进行训练，因此我们挑选ResNet-18这个网络。

• 框架（framework）工具：

这个就是作为深度学习的模型训练与推理识别的重要工具，目前最主流的应该是Tensorflow，不过这里为了配合迁移学习（transfer learning）的特性，因此选择用pytorch这个轻量级的框架工具

• 训练用的设备：

事实上训练用的设备，会直接影响到网络模型与框架工具的选择。这里我们使用NVIDIA Jetson系列设备进行训练，就必须选择轻量级的网络模型与框架工具。

现在就开始我们的图像分类模型训练实验，在开始实验之前，还是有以下几个准备动作需要确认：

• 是否已经安装好pytorch这个工具？如果前面没有装好，现在还可以执行以下指令进行安装：

• 是否为系统添加SWAP空间？如果是正常用Jetack 4.4以上版本的话，都应该已经设定好SWAP空间，如果确认未设定SWAP的话，请执行以下指令：

• 关闭jetson设备的图像桌面，可以腾出更多内存，有利于训练作业。

接下来就正式进入模型训练的环节，项目作者已经为我们提供基于pytorch的图像分类的Python训练代码，为了便于操作，先将工作目录移到代码的位置：

里面的train.py与onnx_export.py是这个训练过程会用到的两只程序，至于代码内容则不作说明。

接下去的任务就是获取数据集。项目作者提供了两个现成的数据集“cat_dog”与“PlantCLEF_Subset”，除了数据集本身的内容（含类别）不一样之外，其余操作过程是完全一样的，因此这里我们只演示“cat_dog”这个数据集的完整过程，将“PlantCLEF_Subset”的实验留给读者去比照办理。

为了保持指令参数的一致性，建议将数据下载到data目录下，并且将生成的模型文件制定到models路径，现在就开始执行完整的步骤：

检查看看cat_dog与PlantCLEF_Subset下面，是否都存在<test>、<train>、<val>三个目录以及一个labels.txt文件？

在cat_dog下面的test、train、val目录下，应该都有cat与dog子目录，这个子目录名称会与labels.txt内容对应上。根据这样的原则，请检视一下PlantCLEF_Subset数据集，是否也具备相同的特性？这里应该有20种植物名称分类。

接下去回到train.py与onnx_export.py工作目录下，准备开始执行模型训练任务：

由于训练程序使用了迁移学习的技巧，因此在第一次开始训练之前，会先下载合适的预训练模型，然后让这些数据在这个模型基础上进行叠加训练，一方面提高精准度，另一方面也会减少训练时间，这就是“迁移学习”的迷人之处。

系统预设使用resnet-18这个网络，这是比较适合Jetson Nano（含2GB）的轻量级模型，我们也可以使用 --arch= 去指定使用其他网络模型。如果这时候打开jetson-stats监控工具，可以看到GPU持续处于满载状态。

这个cat_dog数据集共有6200张图片，其中训练集5,000张、验证集1,000张、测试集200张，系统预设的epoch数为35，可使用--epoch=去改变epoch数。epoch数量设为多少比较合适？这个没有一定的答案，我们需要观察训练过程中Loss值的变化。

如此的数据量在Jetson Nano 2GB上的训练时间，总共花了285分钟（4小时45分钟）

现在训练完成之后，下一步就是要把这个训练的模型，使用onnx_import.py将pytorch模型文件转换成ONNX格式，以便在后面使用时很轻松地转换成TensorRT加速引擎。

会在models/cat_dog目录下生成“resnet18.onnx”文件，最后以这个文件为基础，调用前面所熟悉的imagenet图像分类指令，来进行推理计算任务，请执行以下代码：

现在使用自己训练的模型来执行imagenet.py，因此需要个别给定--model、--labels的路径，以及--input_blob、--output_blob的内容，其他部分与前面讲解过的用法是一致的，可以用视频、摄像头、文件夹作为输入源，检测一下识别的效果如何？

请自行尝试PlantCLEF_Subset这个关于植物类别的数据集，按照上面完全一样的做法，自行走过一遍，这会让你更加熟悉与掌握整个“图像分类”的应用，从数据集收集、模型训练，到推理识别这三个阶段。

最后，如果您需要建立自己专属的数据集的话，这个Hello AI World项目还提供一个camera-capture这个现场抓图的工具，开启之后会出现一个控制面板（如下图左），可选择Classification或Detection功能，同时还会启动一个摄像头画面市场（如下图右），可以自行调整所需要的的分类与指定路径。

收集完足够的图像数据之后，再按照前面的步骤进行模型训练，就能建立自己专属的模型文件，给自己的Jetson设备进行图像分类的推理识别应用。