Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

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

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

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

  • 数据集(dataset:

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

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

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

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

  • 框架(framework)工具:

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

  • 训练用的设备:

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

 

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

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

1

2

  • cd ~/jetson-inference/build
  • ./install-pytorch.sh
  • 是否为系统添加SWAP空间?如果是正常用Jetack 4.4以上版本的话,都应该已经设定好SWAP空间,如果确认未设定SWAP的话,请执行以下指令:

1

2

3

4

  • sudo systemctl disable nvzramconfig
  • sudo fallocate -l 4G /mnt/4GB.swap
  • sudo mkswap /mnt/4GB.swap
  • sudo swapon /mnt/4GB.swap

将“/mnt/4GB.swap  none  swap  sw 0  0”加入 /etc/fstab 文件后重启

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

1

2

3

4

  • sudo init 3          # 临时关闭桌面
  • sudo init 5          # 重新开启桌面
  • sudo systemctl set-default multi-user.target  # 永久关闭桌面,需重启
  • sudo systemctl set-default graphical.target   # 永久开启桌面,需重启

 

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

1

  • cd ~/jetson-inference/python/training/classification

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

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

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

1

 

2

 

 

3

 

4

 

 

5

  • cd data

# 下载并解压缩 cat_dog.tar.gz758MB

  • wget -c  https://nvidia.box.com/shared/static/o577zd8yp3lmxf5zhm38svrbrv45am3y.gz -O cat_dog.tar.gz
  • tar xvzf cat_dog.tar.gz

# 下载并解压缩 PlantCLEF_Subset.tar.gz1.34GB

  • wget -c https://nvidia.box.com/shared/static/vbsywpw5iqy7r38j78xs0ctalg7jrg79.gz -O PlantCLEF_Subset.tar.gz
  • tar xvzf PlantCLEF_Subset.tar.gz

 

检查看看cat_dogPlantCLEF_Subset下面,是否都存在<test><train><val>三个目录以及一个labels.txt文件?

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

接下去回到train.pyonnx_export.py工作目录下,准备开始执行模型训练任务:

 

1

2

  • cd ~~/jetson-inference/python/training/classification/
  • python3 train.py --model-dir=models/cat_dog data/cat_dog

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

Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

1

  • python3 onnx_export.py --model-dir=models/cat_dog

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

1

2

3

  • export NET=models/cat_dog
  • export DATASET=data/cat_dog
  • imagenet.py $DATASET/test/cat/01.jpg classified_01.jpg \   #输入与输出

--model=$NET/resnet18.onnx --labels=$DATASET/labels.txt \ #模型的文件

--input_blob=input_0 --output_blob=output_0

 

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

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

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

Jetson Nano 2GB 系列文章(25): “Hello AI World”图像分类的模型训练

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

来源:业界供稿

0赞

好文章,需要你的鼓励

2021

07/02

15:00

分享

点赞

邮件订阅
白皮书