PyTorchでSSDを試してみた(2) - みらいテックラボ

普段物体検知を行うとき, これまでYOLOv3[1]やssd_keras[2][3]を使ってきた.
しかし, これらの物体検知をJetson Nanoなど組み込みボードで処理しようとすると, フレーム毎の処理に結構時間がかかり, 秒数フレームほどしか処理できなかったりした.
JetsonでDNN処理の高速化を考えた場合, Pythonの実装を頑張って高速化を目指すより, NVIDIAのTensorRTを活用する方が容易そうだ.
そこで, PyTorchのSSD(Single Shot MultiBox Detector)モデルをPCで学習し, 学習済みモデルをONNX形式に変換してJetsonで活用することを試してみることにした.

関連記事：
・PyTorchでSSDを試してみた(1)
・PyTorchでSSDを試してみた(2)
・PyTorchでSSDを試してみた(3)
・PyTorchでSSDを試してみた(4)

f:id:moonlight-aska:20210914221840j:plain:w400

前回は, オリジナルデータでPyTorchのSSDモデルを転移学習する方法について紹介した.
今回は, 学習済みモデルをONNX形式に変換して, jetson-inference[4]の物体検知で使用する方法について紹介していく.

1. Jetson環境の構築
PCで学習したSSDモデルをJetsonで使用するには, JetsonでSSDモデルの変換作業を行う必要がある.
そこで, まずはJetson側の環境を準備する.

1.1 Jetson Nanoセットアップ
Jetson Nanoのセットアップについては, 公式サイト[5]の「Getting Started with Jetson Nano Developer Kit」を参照すれば, 特に迷うことなくセットアップできる.
ただし, Jetson Nano 2GBの方は, 「Getting Started with Jetson Nano 2GB Developer Kit」を参照ください.

ボード	NVIDIA Jetson Nano
RAM	4GB
OSパッケージ	JetPack 4.6

1.2 jetson-inferenceインストール
Jetson Nanoで学習済みモデルのONNX変換や物体検知を試すのに必要なjetson-inferenceをGitHubからダウンロードし, ビルドする.
インストール手順を下記に記載するが, 詳細はこちらを参照.

$ sudo apt-get update
$ sudo apt-get install git cmake libpython3-dev python3-numpy
$ git clone --recursive https://github.com/dusty-nv/jetson-inference
$ cd jetson-inference
$ mkdir build
$ cd build
$ cmake ../
$ make -j$(nproc)
$ sudo make install
$ sudo ldconfig

2. ONNX形式モデル変換
2.1 学習済みモデルの準備
PCで学習したpytorch-ssdの学習済みモデルのうち, Lossが最も小さいpthファイルをJetsonにコピーする.
Lossが小さいpthファイルを探すのが面倒な場合は, すべてのpthファイルをコピーしてもよい.
それから, 物体検知時に必要となるラベルファイルもJetsonにコピーしておく.

$ cd jetson-inference/python/training/detection/ssd
$ scp username@<PCのIPaddress>:<model-dir>/mb1-ssd-Epoch-*.pth models
$ scp username@<PCのIPaddress>:<model-dir>/labels.txt models

2.1 モデル変換
PyTorchのSSDモデルをONNX形式に変換するには, jetson-inferenceのssd ディレクトリの中に含まれているonnx_export.pyを使用する.
モデル変換は簡単で, 学習時のモデルを保存しているディレクトリを指定するだけで, Loss最小のモデルをONNX形式のモデルに変換してくれる.

$ python3 onnx_export.py --model-dir=models
Namespace(batch_size=1, height=300, input='', labels='labels.txt', model_dir='models', net='ssd-mobilenet', output='', width=300)
running on device cuda:0
found best checkpoint with loss 1.737015 (models/mb1-ssd-Epoch-48-Loss-1.7370146930217742.pth)
creating network:  ssd-mobilenet
num classes:       2
loading checkpoint:  models/mb1-ssd-Epoch-48-Loss-1.7370146930217742.pth
exporting model to ONNX...
graph(%input_0 : Float(1:270000, 3:90000, 300:300, 300:1),
      %base_net.0.0.weight : Float(32:27, 3:9, 3:3, 3:1),
      %base_net.0.1.weight : Float(32:1),
      %base_net.0.1.bias : Float(32:1),
         :
         :
  %469 : Float(1:6000, 3000:2, 2:1) = onnx::Div(%467, %468)
  %470 : Float(1:6000, 3000:2, 2:1) = onnx::Add(%466, %469) # /home/aska/work/pytorch-ssd/vision/utils/box_utils.py:209:0
  %boxes : Float(1:12000, 3000:4, 4:1) = onnx::Concat[axis=2](%465, %470) # /home/aska/work/pytorch-ssd/vision/utils/box_utils.py:209:0
  return (%scores, %boxes)

model exported to:  models/ssd-mobilenet.onnx
task done, exiting program

2.2 動作確認
学習済みモデルを試すコードを作成し, モデルを試してみた.
jetson-inferenceの物体検知を使用するにあたり, いくつか注意すべき点がある. (コード参照)
・画像のカラーモードはRGBAである
・画像データはnumpyのndarrayからGPUメモリへコピーして渡す.
・detectNetはコマンドラインオプション形式でカスタムモデルなどの引数を指定する.

物体検知ブログラムの初回実行時には, onnx形式のモデルをFP16に変換するようで, かなり時間がかかる.
初回実行後には, modelsの下に, "ssd-mobilenet.onnx.1.1.8001.GPU.FP16.engine"というファイルができていた.
[コード]

import jetson.inference
import jetson.utils
import numpy as np
import cv2

network = "ssd-mobilenet-v1"
argv = ["--model=models/ssd-mobilenet.onnx", 
                "--labels=models/labels.txt",
                "--input-blob=input_0",
                "--output-cvg=scores",
                "--output-bbox=boxes" ]
net = jetson.inference.detectNet(network, argv, threshold=0.5)

bgr_img = cv2.imread('./data/FCam9_00034702.jpg')
h, w, _ = bgr_img.shape
rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2RGBA).astype(np.float32)
img = jetson.utils.cudaFromNumpy(img)
detections = net.Detect(img)
img = jetson.utils.cudaToNumpy(img)
rgb_img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB).astype(np.uint8)
bgr_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2BGR)
cv2.imwrite('output.jpg', bgr_img)
cv2.imshow('SSD Win', bgr_img)
cv2.waitKey()
cv2.destroyAllWindows()

[実行結果]
f:id:moonlight-aska:20210912234747j:plain:w400