みらいテックラボ

音声・画像認識や機械学習など, 管理人が興味のある技術の紹介や実際にトライしてみた様子などメモしていく.

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)


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

他のデータの結果
f:id:moonlight-aska:20210912234840j:plain:w400

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

今回は, PCで学習したSSDモデルをONNX形式に変換し, jetson-inferenceの物体検知で利用する方法について紹介した.
次回は少しフレーム毎の処理について, 実行時間などを確認してみようと思う.


----
[1] YOLO: Real-Time Object Detection
[2] GitHub - pierluigiferrari/ssd_keras: A Keras port of Single Shot MultiBox Detector
[3] GitHub - rykov8/ssd_keras: Port of Single Shot MultiBox Detector to Keras
[4] GitHub - dusty-nv/jetson-inference: Hello AI World guide to deploying deep-learning...
[5] Jetson Nano Developer Kit | NVIDIA Developer
[6] jetson-inference/building-repo-2.md at master・dusty-nv/...







PyTorchではじめるAI開発

PyTorchではじめるAI開発

Amazon