みらいテックラボ

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

PyTorchでSSDを試してみた(1)

普段物体検知を行うとき, これまで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を使って撮影した距離画像を使用した.
[画像のサンプル]
f:id:moonlight-aska:20210911224715p:plain

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

[結果]
f:id:moonlight-aska:20210911235629j:plain:w400

次回は, この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...







PyTorchではじめるAI開発

PyTorchではじめるAI開発

Amazon