普段物体検知を行うとき, これまで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)
1. 開発環境の準備
1.1 開発環境
今回の開発環境は以下の通り.
| CPU | Intel Core i7-7700 |
| RAM | 16GB |
| OS | Ubuntu 18.04 |
| GPU | GeForce GTX 1060 6GB |
[GPU環境]
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.33.01 Driver Version: 440.33.01 CUDA Version: 10.2 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 GeForce GTX 106... On | 00000000:01:00.0 Off | N/A |
| 38% 30C P8 5W / 120W | 1MiB / 6078MiB | 0% Default |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: GPU Memory |
| GPU PID Type Process name Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+1.2 pytorch-ssdの準備
PyTorchでSSDを試すにあたり, ベースとなりそうなものはいくつかあるが, 最終的にはJetson Nanoで動かしたいので, jetson-inference[4]のpytorch-ssd[5]をベースに試してみた.
$ git clone https://github.com/dusty-nv/pytorch-ssd $ cd pytorch-ssd $ wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/mobilenet-v1-ssd-mp-0_675.pth $ pip3 install -v -r requirements.txt
2. 学習データの準備
2. 1 学習データ
学習データには, 以前ピープルカウンタ開発[6]で使用したIntel Real Senseを使って撮影した距離画像を使用した.
[画像のサンプル]
2.2 アノテーションデータ
以前に, マイクロソフト社のアノテーションツールVoTT[7]を使って作業をして, Pascal VOC形式で出力したものを使用した.
[アノテーションデータの構造]
<annotation verified="yes">
<folder>Annotation</folder>
<filename>FCam9_00036372.jpg</filename>
<path>Hug2-gray-2018-PascalVOC-export/Annotations/FCam9_00036372.jpg</path>
<source>
<database>Unknown</database>
</source>
<size>
<width>640</width>
<height>480</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>Passer</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>292.54571026722925</xmin>
<ymin>251.81988742964353</ymin>
<xmax>418.56540084388183</xmax>
<ymax>348.18011257035647</ymax>
</bndbox>
</object><object>
<name>Passer</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>501.3783403656821</xmin>
<ymin>245.51594746716697</ymin>
<xmax>611.1954992967651</xmax>
<ymax>344.577861163227</ymax>
</bndbox>
</object>
</annotation>2.3 データの配置
今回train_ssd.pyで学習するにあたり, データは"pytorch-ssd/data"の下に, 以下のように配置した.
[データ配置構造]
data/Hug2-gray-2018-PascalVOC-export ├── Annotations │ ├── FCam9_00000154.xml │ ├── FCam9_00000156.xml │ ├── FCam9_00000158.xml │ ├── FCam9_00000160.xml │ ├── FCam9_00000162.xml │ ├── (省略) ├── ImageSets │ └── Main │ ├── test.txt │ ├── trainval.txt ├── JPEGImages │ ├── FCam9_00000154.jpg │ ├── FCam9_00000156.jpg │ ├── FCam9_00000158.jpg │ ├── FCam9_00000160.jpg │ ├── FCam9_00000162.jpg │ ├── (省略) ├── labels.txt
3. 転移学習
3.1 モデル学習
上記の学習データを使って, train_ssd.pyで転移学習を行う.
train_ssd.pyのパラメータの一部を紹介する.
| パラメータ | 説明 |
| --dataset-type | データセットのタイプを指定する. "voc"と"open_images"のいずれか. デフォルトは"open_images" |
| --datasets or --data | データセットのディレクトリパスを指定する. デフォルトは"data" |
| --net | ネットワークアーキテクチャを指定する. "vgg16-ssd", "mb1-ssd", "mb1-ssd-lite", "mb2-ssd-lite", "sq-ssd-lite". デフォルトは"mb1-ssd" |
| --pretrained-ssd | 学習済みモデルを指定する. デフォルトは"models/mobilenet-v1-ssd-mp-0_675.pth" |
| --resume | 学習を再開するためのチェックポイントへのパスを指定する. |
| --batch_size | ミニバッチのサイズを指定する. デフォルトは4. |
| --num_epochs or --epochs | エポック数を指定する. デフォルトは30. |
| --num-workers or --workers | スレッド数を指定する. デフォルトは2 |
| --checkpoint-folder or --model-dir | モデルのチェックポイントを保存するディレクトリを指定する. デフォルトは"models/" |
$ python3 train_ssd.py --dataset-type=voc --data=data/Hug2-gray --model-dir=models/Hug2 --batch-size=32 --num-workers=4 --epochs=50
2021-09-11 23:26:34 - Using CUDA...
2021-09-11 23:26:34 - Namespace(balance_data=False, base_net=None, base_net_lr=0.001, batch_size=32, checkpoint_folder='models/Hug2', dataset_type='voc', datasets=['data/Hug2-gray'], debug_steps=10, extra_layers_lr=None, freeze_base_net=False, freeze_net=False, gamma=0.1, lr=0.01, mb2_width_mult=1.0, milestones='80,100', momentum=0.9, net='mb1-ssd', num_epochs=50, num_workers=4, pretrained_ssd='models/mobilenet-v1-ssd-mp-0_675.pth', resume=None, scheduler='cosine', t_max=100, use_cuda=True, validation_epochs=1, weight_decay=0.0005)
2021-09-11 23:26:34 - Prepare training datasets.
False data/Hug2-gray/ImageSets/Main/trainval.txt
2021-09-11 23:26:35 - VOC Labels read from file: ('BACKGROUND', 'Passer')
2021-09-11 23:26:35 - Stored labels into file models/Hug2/labels.txt.
2021-09-11 23:26:35 - Train dataset size: 2528
2021-09-11 23:26:35 - Prepare Validation datasets.
True data/Hug2-gray/ImageSets/Main/test.txt
2021-09-11 23:26:35 - VOC Labels read from file: ('BACKGROUND', 'Passer')
2021-09-11 23:26:35 - Validation dataset size: 632
2021-09-11 23:26:35 - Build network.
2021-09-11 23:26:35 - Init from pretrained ssd models/mobilenet-v1-ssd-mp-0_675.pth
2021-09-11 23:26:35 - Took 0.03 seconds to load the model.
（省略)
2021-09-11 23:50:01 - Epoch: 49, Step: 50/79, Avg Loss: 1.9336, Avg Regression Loss 0.5225, Avg Classification Loss: 1.4111
2021-09-11 23:50:04 - Epoch: 49, Step: 60/79, Avg Loss: 1.8585, Avg Regression Loss 0.5216, Avg Classification Loss: 1.3369
2021-09-11 23:50:07 - Epoch: 49, Step: 70/79, Avg Loss: 1.8944, Avg Regression Loss 0.5193, Avg Classification Loss: 1.3752
2021-09-11 23:50:12 - Epoch: 49, Validation Loss: 2.1572, Validation Regression Loss 0.6638, Validation Classification Loss: 1.4934
2021-09-11 23:50:12 - Saved model models/Hug2/mb1-ssd-Epoch-49-Loss-2.1571941018104552.pth
2021-09-11 23:50:12 - Task done, exiting program.3.2 動作確認
run_ssd_example.pyを使って, 学習したモデルの動作検証を行う.
$ python3 run_ssd_example.py mb1-ssd model/Hug2/mb1-ssd-Epoch-49-Loss-2.1571941018104552.pth models/Hug2/labels.txt data/FCam9_00036372.jpg (480, 640, 3) Inference time: 0.005816221237182617 Found 2 objects. The output image is run_ssd_example_output.jpg
[結果]
次回は, このPyTorch 形式のモデル(*.pth)をNVIDIA TensorRT でも扱えるONNX(Open Neural Network Exchange)形式に変換して, jetson-inferenceのdetectnet.pyで試してみようと思う.
----
[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] GitHub - dusty-nv/pytorch-ssd: MobileNetV1, MobileNetV2, VGG based SSD/SSD-lite ...
[6] ピープルカウンタを考えてみる(1)～(8)
[7] GitHub - microsoft/VoTT: Visual Object Tagging Tool: An electron app for building end...
