训练图像分类模型

2023-07-25 admin

训练图像分类模型

1.训练前的准备

在这个教程开始前，python必须安装好torchCPU+GPU版本，可以看本系列教程的torch的安装教程。

2.开始对猫狗数据集重新训练(我司提供的压缩包里包含此项,无需执行此步骤)

cd jetson-inference/python/training/classification/data

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

tar xvzf cat_dog.tar.gz

如果图片因为网络问题下载不成功，可以在PC上的"模型/cat_dog"中找到cat_dog.tar.gz传输到NX上;

3.执行训练的命令

cd jetson-inference/python/training/classification

python3 train.py --model-dir=models/cat_dog data/cat_dog

注意:如果内存不足导致训连过程中被“终止”,请尝试交换虚拟内存和禁用桌面GUI.

#禁用桌面GUI

sudo init 3 # stop the desktop

sudo init 5 # restart the desktop

交换虚拟内存（参考链接）：https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-transfer-learning.md#mounting-swap

4.训练开始的信息解读

Use GPU: 0 for training

=> dataset classes: 2 ['cat', 'dog']

=> using pre-trained model 'resnet18'

=> reshaped ResNet fully-connected layer with: Linear(in_features=512, out_features=2, bias=True)

Epoch: [0][ 0/625] Time 0.932 ( 0.932) Data 0.148 ( 0.148) Loss 6.8126e-01 (6.8126e-01) Acc@1 50.00 ( 50.00) Acc@5 100.00 (100.00)

Epoch: [0][ 10/625] Time 0.085 ( 0.163) Data 0.000 ( 0.019) Loss 2.3263e+01 (2.1190e+01) Acc@1 25.00 ( 55.68) Acc@5 100.00 (100.00)

Epoch: [0][ 20/625] Time 0.079 ( 0.126) Data 0.000 ( 0.013) Loss 1.5674e+00 (1.8448e+01) Acc@1 62.50 ( 52.38) Acc@5 100.00 (100.00)

Epoch: [0][ 30/625] Time 0.127 ( 0.114) Data 0.000 ( 0.011) Loss 1.7583e+00 (1.5975e+01) Acc@1 25.00 ( 52.02) Acc@5 100.00 (100.00)

Epoch: [0][ 40/625] Time 0.118 ( 0.116) Data 0.000 ( 0.010) Loss 5.4494e+00 (1.2934e+01)

Epoch: [0]:代表训练到第几个周期 0：第一个如果没指定参数，默认是只训练35个周期可以使用--epochs=N指定训练多少周期 N=100,即训练100个周期

[N/625]:是您所在时代的当前图像批次小批量处理训练图像以提高性能,默认的批处理大小是8个图像，可以用--batch=N 将括号中的数字乘以批次大小(例如批次[100/625]->图像[800/5000]) -Time:当前图像批次的处理时间(秒) -Data:当前图像批次的磁盘加载时间(秒) -Loss:模型产生的累积误差(预期误差和预测误差) 如果你想训练10次，批处理的图像为6，即运行一下命令

python3 train.py --model-dir=models/cat_dog data/cat_dog --epochs=10 --batch=6

5.模型精度讲解

在这个由5000幅图像组成的数据集上，在Jetson orin NX上训练ResNet-18每个时期大约需要2分钟，或者大约55分钟来训练模型达到35个时期和80%的分类精度。下图用于分析历元的训练进度与模型准确性的关系:

在大约第30个时期，ResNet-18模型达到80%的精确度，在第65个时期，它收敛于82.5%的精确度。有了额外的训练时间，您可以通过增加数据集的大小来进一步提高准确性或者尝试更复杂的模型。默认情况下，训练脚本设置为运行35个时期，但是如果您不希望等待那么长时间来测试您的模型，您可以提前退出训练并继续下一步(可以选择稍后从您停止的地方重新开始训练)。通过--resume这个参数进行恢复.

6.onnx模型转换

onnx的模型文件方便TensorRT进行处理

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

这将创建一个名为resnet18.onnx在jetson-inference/python/training/classification/models/cat_dog/下面

7.用TensorRT处理图像

举例

cd /home/jetson/jetson-inference/python/training/classification

export NET=models/cat_dog

export DATASET=data/cat_dog

imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt $DATASET/test/cat/01.jpg cat.jpg

以上的命令执行完毕，可以在下面的路径观察的到结果，该命令是处理一张图片路径：/home/jetson/jetson-inference/python/training/classification

如果你想要测试所有的图片可以输入一下命令

cd /home/jetson/jetson-inference/python/training/classification

export NET=models/cat_dog

export DATASET=data/cat_dog

imagenet --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/../labels.txt \

$DATASET/test/dog ./data/cat_dog/test_output_dog

输出的结果在/home/jetson/jetson-inference/python/training/classification/data/cat_dog/test_output_dog路径下

如果你想要打开摄像头进行检测

imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt csi://0 #对CSI摄像头

imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt v4l2:///dev/video0 #对v4l2摄像头

其它视频流文件的打开方式，可参考：https://github.com/dusty-nv/jetson-inference/blob/master/docs/aux-streaming.md

8.训练结果

1个训练周期

35个训练周期

35个训练周期识别结果

100个训练周期的结果

可以看到100个训练周期的结果比35个训练周期的结果的精度大大提升。

9.附录

如果你想对自己的各种类别不同的图片进行训练可以参考以下的链接： https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-collect.md

Jetson 共用资料

训练图像分类模型