前一篇我們玩了數種不同的影像分類模型，都是以 ImageNet 為資料集訓練好的預訓練模型，但如果我們想要訓練自己的影像分類模型呢？今天就來訓練自己的影像分類模型吧！為了方便起見，我們將直接在 Jetson Nano 上進行模型訓練！

Cat or Dog? (source: NVIDIA)

下載資料集

訓練模型不需要進到 docker 容器中，直接在 Jetson Nano 操作即可。首先需要下載資料集，這次貓狗的影像分類資料集由 NVIDIA 提供，並且已做好分類，每一類別資料為訓練資料(train) 2500 張，驗證資料(validation) 500 張與測試資料(test) 100 張。資料集不算大，但也有 723MB，完成下載約須 3 分鐘。下載完直接進行解壓縮。

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

重新訓練 ResNet-18 模型

切換資料夾至訓練的工作目錄，輸入 train.py 開始進行訓練：

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

預設一樣使用 ResNet-18 預訓練模型做遷移學習，因此會先下載 ResNet-18 權重，但是因為資料量較大，每一次訓練迭代時間較久，估計 7 分鐘左右才會完成一次訓練。過程可以看到訓練過程中的 log 輸出資訊：

• Epoch: [N] - 顯示第 N 次的訓練，預設為訓練35次，可以使用 --epochs=N 參數來指定訓練次數。
• [N/625] - 當前批次訓練進度，預設 batch size 為 8，因為總共有 5000張訓練資料（每類2500張*2類），每次隨機抓取8張影像，所以每次訓練會有625次迭代，若batch size 有調整，這裡的分母 625 也會變動。 可用 --batch=N 指定 batch size。
• Loss - 推論誤差，此數值越低越好。
• Acc@1: 最高機率命中正確類別的準確度。
• Acc@5: 前五高機率命中正卻類別的準確度。（因為資料只有兩個類別，所以 Top-5 一定是100%命中 XD）

訓練地獄就此開始！這邊可以先去洗個澡，叫個外賣，吃個飯，上個廁所…..

過程如果真的受不了等待，也可以按 Ctrl + C 中斷，程式會保留目前為止最佳的推論模型。抑或可以直接下載官方在 100 次訓練後的模型來執行後續測試。

轉換模型為 ONNX 格式

執行下方指令將 pytorch 的 pth 模型檔轉換為 ONNX 格式，完成後會在原資料夾多一個名為 resnet18.onnx 模型。

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

使用 TensorRT 加速推論

接著如同前兩篇一樣，我們先進到 docker 容器中進行推論：

cd ~/jetson-inference/
docker/run.sh

執行以下指令指定測試照片進行推論：

NET=/jetson-inference/python/training/classification/models/cat_dog
DATASET=/jetson-inference/python/training/classification/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

推論後的結果僅有 50% 的機率是貓（好在也是過半了！），這也是因為筆者測試的模型也僅只跑 21 次的訓練即中斷了，若把訓練次數增加應會有效提高正確性。除此之外，訓練好的模型也可以如同上一篇的說明，使用 camera 即時推論或是做批量照片的推論等，這邊就不除重複贅述了，有興趣的開發者可以自行把玩！

小結

即便 Jetson Nano 的 GPU 也是可以訓練模型的，但實際跑起來確實是有非常大的無力感！但這是為了方便起見，不用在 PC 端去建置另一個開發環境，做個實驗性質的測試並體驗訓練模型的過程。正常合理情況下勢必會是在個人工作站或是伺服器上做模型的訓練啦！