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

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

1. 数据集（dataset）:

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

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

网上已有的开源资源：包括ImageNet、MS COCO、OpenImages 等等，都是数量级在百万张图片以上的数据集，而且分类十分清楚，是个非常重要的资源，必须学会如何从这些数据集中，萃取自己所需要的的数据。

自行收集：通过在网上根据关键字去搜索下载，也可以利用自己的摄像设备去拍摄收集。

自行转换：可以将已有的图像数据，经过灰度处理、镜像映射等方式，再生成 2x~4x 数据量。

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

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

3. 框架（framework）工具：

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

4. 训练用的设备：

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

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

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

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

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

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

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

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

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

检查看看 cat_dog 与 PlantCLEF_Subset 下面，是否都存在 、、 三个目录以及一个 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 设备进行图像分类的推理识别应用。