PyTorch SSDで転移学習（mobilenetベースSSD）

以下記事で、VGG16ベースのSSDの転移学習を実施しましたが、SSDのベースネットワークをVGG16→mobilenet(v2-lite)に変えて、転移学習を実施しました。

https://qiita.com/shima202102/items/a92a28ff277bcea9202f

https://qiita.com/shima202102/items/4e24d77df70448acae1f

実装は、以下githubに公開してますので、ダウンロードしてご利用ください。

https://github.com/Shimasaki2021/pytorch_ssd_trial

# 背景

以下記事で、VGG16ベースのSSDに対して転移学習を行い、ナンバープレートぼかしアプリを作成したのですが、SSDの検知が遅く（※）、約30分ぐらいの動画全てを処理するのに7.5時間かかってしまう有様でした。

この処理時間の課題をなんとかしたいと思い、今回、mobilenetベースの転移学習に挑戦してみました。

https://qiita.com/shima202102/items/0bea767122c288021194

（※）下記PC環境で2fps程度。1分の動画を処理するのに14,5分ほどかかる。

- PC環境
  - CPU: AMD Ryzen 7 3700X (3.60 GHz)
  - GPU: NVIDIA GeForce GTX 1660 SUPER
  - OS: Windows 11 Home （24H2） , WSL2 + Ubuntu24.04


# 目次

作成したいモデル仕様（検知対象＝車＆ナンバープレートetc.）や、pytorch環境構築は、VGG16ベースのSSDと共通です。以前の記事（以下）を参照願います。

https://qiita.com/shima202102/items/a92a28ff277bcea9202f

1. SSD実装(mobilenetベースSSD移植)
2. 転移学習実行
3. 検知（推論）実行、評価

# 1. SSD実装(mobilenetベースSSD移植)

VGGベースSSDの転移学習時に作成したソース一式（```pytorch_ssd_trial```（以下github））に、Hao Gao様のgithub（Copyright (c) 2019 Hao Gao）で公開されているmobilenetベースSSD実装を移植しました。

https://github.com/Shimasaki2021/pytorch_ssd_trial

https://github.com/qfgaohao/pytorch-ssd

ここには、色々なバージョンのSSDが実装されています。mobilenet(v1,v2,v3)ベースのSSDだけでなく、定番のVGG16ベースのSSDもありました。

この中から、mobilenet-v2-liteベースのSSD実装を選択しました。mobilenet-v3が一番新しいのですが、v2と比較して精度は上がるが処理時間がややかかる、という評価のようなので、最速のv2-liteを選択しました。

移植した箇所は、下図赤枠です。
![fig_soft_structure_mobilenet.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1299377/f2f7a8b4-4734-499b-9726-295f98fbd9d3.png)

Hao Gao様のgithubから、上図赤枠内のmobilenet-v2-lite-ssdのモデル実装ソースのみ拝借し、データローダーや前処理、損失関数の実装は、既存のものをそのまま流用する方針で移植を行いました。

移植で必要となった主な作業は以下です。
- 順伝播のI/F適合
- DefaultBoxサイズ調整

## 1.1 順伝播のI/F適合

順伝播は、以下I/Fの相違があったので、それを、移植先の既存実装（```pytorch_ssd_trial```）に合わせました。

### 学習時の順伝播（```SSD.forward()```）

学習時の順伝播（```SSD.forward()```）のI/F（入力、出力）の相違は下表です。

|  | 入力 | 出力 |
|---|---|---|
| 移植元 | 前処理後の複数枚(batch)画像(```torch.Tensor```) | batch数分のオフセット情報loc、信頼度conf |
| 移植先（```pytorch_ssd_trial```） | ※同上（前処理後の複数枚(batch)画像(```torch.Tensor```)） | batch数分のオフセット情報loc、信頼度conf、DefaultBox |

入力は同じです。
出力の相違は、```DefaultBox```の有無だけです。なので、```SSD.forward()```の出力に```DefaultBox```を追加するだけの変更になってます。

![fig_diff_ssd_ssd.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1299377/22cd6b2d-ffe9-4f49-95ee-ea6bf2188ca5.png)

移植元と移植先で、DefaltBoxの変数名が違っていたので（```（移植元）priors, （移植先）dbox_list```）、対応を取るのが難しかったですが、それさえわかれば、後は追記するだけなので楽でした。

あと、学習時に、mobilenetベースSSD固有の関連として、本節冒頭のクラス図には、```MultiboxLoss(損失関数)```と```AdamW(Optimizer)```からも線（緑の線）が延びてます。なので、ここにも手を入れる必要があるように思えてしまいますが、ここへの手入れは不要でした。

特に、```MultiboxLoss(損失関数)```は、SSDモデルの勾配パラメータとどうつながっているのかが、pytorchの外側からは全く見えないので、手の入れようがありませんでした。とはいえ、手を入れなくてもよさげに動いているので、現状、良しとしてます。

これは、今回のmobilenetベースSSD固有ではなく、前回のVGG16ベースSSDのコードでも理解できていない部分です。ご存じの方がいらっしゃいましたら、教えて下さると助かります。


### 検知時の順伝播（```Predictor.predict()```）

検知時の順伝播（```Predictor.predict()```）のI/F（入力、出力）の相違は下表です。

|  | 入力 | 出力 |
|---|---|---|
| 移植元 | 前処理前の画像1枚(```numpy.ndarray```) |  batch数分の信頼度conf、クラス番号(ラベル),　BoundingBox(bbox) |
| 移植先（```pytorch_ssd_trial```） | 前処理後の複数枚(batch)画像(```torch.Tensor```) | batch数＆クラス数分の信頼度conf、BoundingBox(bbox) |

こちらは、入力、出力共に、かなり相違があります。

出力は、どちらも実質同じ情報を返しており、上表からは違いがわからないですが、実装レベルで、返り値の型に以下相違があります。

- 移植元： ```Tuple[torch.tensor[batch,4(bbox)],  torch.tensor[batch](class),  torch.tensor[batch](conf)]```
- 移植先： ```torch.tensor[batch,class,top_k,5(※)]``` (※)5: ```[conf, 4(bbox)]```

ソースは以下です。

既実装の```utils.ssd_model.Detect.forward()```からソースを引っ張ってきて、移植元の大半を置き換えました。ついでに、呼び出し元のソースを書き換えなくてもいいよう、メソッド名も```forward() → __call__()```に変更しました。

```python: predictor.py
# 移植元(Predict.forward())
def predict(self, image, top_k=-1, prob_threshold=None):
    cpu_device = torch.device("cpu")
    height, width, _ = image.shape
    image = self.transform(image)
    images = image.unsqueeze(0)
    images = images.to(self.device)
    with torch.no_grad():
        self.timer.start()
        scores, boxes = self.net.forward(images)
        print("Inference time: ", self.timer.end())
    boxes = boxes[0]
    scores = scores[0]
    if not prob_threshold:
        prob_threshold = self.filter_threshold
    # this version of nms is slower on GPU, so we move data to CPU.
    boxes = boxes.to(cpu_device)
    scores = scores.to(cpu_device)
    picked_box_probs = []
    picked_labels = []
    for class_index in range(1, scores.size(1)):
        probs = scores[:, class_index]
        mask = probs > prob_threshold
        probs = probs[mask]
        if probs.size(0) == 0:
            continue
        subset_boxes = boxes[mask, :]
        box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
        box_probs = box_utils.nms(box_probs, self.nms_method,
                                  score_threshold=prob_threshold,
                                  iou_threshold=self.iou_threshold,
                                  sigma=self.sigma,
                                  top_k=top_k,
                                  candidate_size=self.candidate_size)
        picked_box_probs.append(box_probs)
        picked_labels.extend([class_index] * box_probs.size(0))
    if not picked_box_probs:
        return torch.tensor([]), torch.tensor([]), torch.tensor([])
    picked_box_probs = torch.cat(picked_box_probs)
    picked_box_probs[:, 0] *= width
    picked_box_probs[:, 1] *= height
    picked_box_probs[:, 2] *= width
    picked_box_probs[:, 3] *= height
    return picked_box_probs[:, :4], torch.tensor(picked_labels), picked_box_probs[:, 4]
```

```python: predictor.py
# 移植先(Predict.__call__())
def __call__(self, images:torch.Tensor, top_k=200, prob_threshold=None):

    if not prob_threshold:
        prob_threshold = self.filter_threshold

    with torch.no_grad():
        # self.timer.start()
        scores, boxes = self.net.forward(images)
        # print("Inference time: ", self.timer.end())
        # boxes:  [num_batch,3000,4]
        # scores: [num_batch,3000,21]

    # CPU実行のほうが高速
    cpu_device = torch.device("cpu")
    boxes  = boxes.to(cpu_device)
    scores = scores.to(cpu_device)

    # 出力の型を作成する。テンソルサイズは[minibatch数, num_classes, top_k, 5(score,bbox[4])]
    num_batch   = images.size(0)
    num_classes = scores.size(2)
    output = torch.zeros(num_batch, num_classes, top_k, 5)

    for batch_idx in range(num_batch):

        for class_index in range(1, num_classes):
            probs = scores[batch_idx, :, class_index]
            mask  = probs > prob_threshold
            probs = probs[mask]

            if probs.size(0) == 0:
                continue

            subset_boxes = boxes[batch_idx, mask, :]

            # 3. Non-Maximum Suppressionを実施し、被っているBBoxを取り除く
            #    ids  ：confの降順にNon-Maximum Suppressionを通過したindexが格納
            #    count：Non-Maximum Suppressionを通過したBBoxの数
            ids, count = nm_suppression(subset_boxes, probs, self.iou_threshold, top_k)

            # outputにNon-Maximum Suppressionを抜けた結果を格納
            output[batch_idx, class_index, :count] = torch.cat((probs[ids[:count]].unsqueeze(1), 
                                                                subset_boxes[ids[:count]]), 1)

    return output
```


## 1.2 DefaultBoxサイズ調整

移植元のmobilenet-v2-lite-ssdでは、DefaultBoxサイズの最小値＝60pixelとなっていたのですが（以下ソース）、これだと、ナンバープレートが全く検知できませんでした。

そこで、ナンバープレートを検知できるよう、DefaultBoxサイズを小さめに調整しました。幸い、以下ソースの数値を変更するだけで対応できました。

```python:mobilenetv1_ssd_config.py
# 変更前
specs = [
    SSDSpec(19, 16, SSDBoxSizes(60, 105), [2, 3]),
    SSDSpec(10, 32, SSDBoxSizes(105, 150), [2, 3]),
    SSDSpec(5, 64, SSDBoxSizes(150, 195), [2, 3]),
    SSDSpec(3, 100, SSDBoxSizes(195, 240), [2, 3]),
    SSDSpec(2, 150, SSDBoxSizes(240, 285), [2, 3]),
    SSDSpec(1, 300, SSDBoxSizes(285, 330), [2, 3])
]

# 変更後
specs = [
    SSDSpec(19, 16, SSDBoxSizes(21, 60), [2, 3]),
    SSDSpec(10, 32, SSDBoxSizes(60, 150), [2, 3]),
    SSDSpec(5, 64, SSDBoxSizes(150, 195), [2, 3]),
    SSDSpec(3, 100, SSDBoxSizes(195, 240), [2, 3]),
    SSDSpec(2, 150, SSDBoxSizes(240, 285), [2, 3]),
    SSDSpec(1, 300, SSDBoxSizes(285, 330), [2, 3])
]
```

上記ソースを見ると、DefaultBoxのサイズ以外にも、特徴マップの解像度も書かれており、最大解像度は```19（＝19x19）```となってます。

ナンバープレートのような小さな物体だと、解像度が粗すぎる（小さすぎる）と特徴がつぶれてしまうので、この解像度も大きくしたかった（せめてVGGベースSSDと同じ「38」）のですが、残念ながら上記ソースの数値を変えるだけだと、内部でテンソルサイズが合わなくなるらしく、エラーになりました。

上述の数値だけでなく、ネットワークの構造も変更する必要がありそうですが、そこを変更できるほどSSD実装を理解できてないので、断念しました。


# 2. 転移学習実行

1章で移植したmobilenetベースSSDのパラメータを学習します。
今回の初期パラメータは、以下を使わせていただきました。

- ``` mb2-imagenet-71_8.pth ```
  - https://drive.google.com/drive/folders/1pKn-RifvJGWiOx0ZCRLtCXM5GT5lAluu

これを上述URLからダウンロードし、```pytorch_ssd_trial/weights/```に格納します。

学習データセットは、データローダーを既存のものを流用したので、[前回のVGGベースSSDの転移学習](https://qiita.com/shima202102/items/4e24d77df70448acae1f)で使用したデータセット(labelImgでアノテーション)をそのまま使えます。

コマンドラインから以下入力すると、学習が開始します。

```sh
# ubuntuターミナルの場合
./train_ssd.py ([epoch数])

# Anaconda Powershell Promptの場合
python train_ssd.py ([epoch数])
```

以下のように、epoch毎にloss値をコンソール出力しながら、学習が進みます。このあたりは、前回のVGGベースの転移学習と全く同じです。
![snap_train_ssd_py.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1299377/7ebdde55-b6be-4f49-b9e8-e55a0cac4374.png)

今回は、レイヤのfreezeは行いませんでした。
どうもVGGのように低レイヤに負荷が集中しているわけでもなさそうで、freezeを行っても学習時間がそれほど短くはなりませんでした。なので、学習時間よりも精度を重視し、freezeはなしとしました。

学習にかかる時間は、PC環境や学習データ数に大きく依存しますが、今回、以下PC環境、学習データ数で、約2時間でした。VGGベースSSDの転移学習時と同じPC環境、学習データセットですが、約1時間早く学習が終わりました。パラメータ数が少ない分、学習も早いようです。

- PC環境
  - CPU: AMD Ryzen 7 3700X (3.60 GHz)
  - GPU: NVIDIA GeForce GTX 1660 SUPER
  - OS: Windows 11 Home (24H2） , WSL2 + Ubuntu24.04

- 学習データセット（※検証用（1割）込み）
  - 画像数: 598
  - 物体数（枠の数）: 1539　（車：922、ナンバープレート：617）

学習中の```train_Loss、val_Loss```のグラフは下図でした。```out weight```は、パラメータ出力時の値です。
![fig_graph_train.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1299377/d67f2bac-01be-44b3-9cf0-64a1fdd85097.jpeg)

# 3. 検知（推論）実行、評価

続いて、検知（推論）です。

コマンドラインから以下入力すると、検知（推論）が実行されます。

```sh
# ubuntuターミナルの場合
./predict_ssd.py [動画(mp4) or 画像ファイルパス]

# Anaconda Powershell Promptの場合
python predict_ssd.py [動画(mp4) or 画像ファイルパス]
```

結果例は以下です（左：mobilenetベースSSD、右：VGGベースSSD）。
![fig_result_predict_py.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1299377/52e68beb-15e3-4eee-86e0-683e6edf3857.png)

今回転移学習したmobilenetベースSSDだと、一番手前の銀色の車とナンバープレートは検知できていますが、やや遠方の黒い車、赤い車が検知できていません。この例だけ見ると、残念ながら、精度面では、VGGベースSSDより劣るようです。

もう少し客観的に比較評価ができるよう、論文等でよく出てくる、mAP（Mean Average Precision）も算出してみました。学習データをそのまま使って算出したので数値が高めです。あくまで、下記2モデルの比較用の数値です。

| モデル | mAP(IoU=0.4) | AP(車) | AP(ナンバープレート) |
|---|---|---|---|
| mobilenetベースSSD | 0.85 | 0.91 | 0.79 |
| VGGベースSSD | 0.90 | 0.96 | 0.84 |

やはり、VGGベースのSSDと比較すると、mobilenetベースSSDは少し低めの数値となっており、精度はやや劣っています。

次に処理速度の比較です。1分の動画を処理したときの、1秒あたりの処理frame数（fps）で比較しました。

| モデル | fps |
|---|---|
| mobilenetベースSSD | 33.8 |
| VGGベースSSD | 5.1 |

こちらは、mobilenetベースSSDが圧倒的に速いです。リアルタイムの処理速度が出ています。

上述の精度と速度が、具体的にどれくらいのものなのか、検知中のウィンドウ表示を動画キャプチャーしてみました。こちらは、ナンバープレート自動ぼかしアプリの処理の様子をキャプチャしています。VGGベースSSDとmobilenetベースSSDを並べてみました。

<iframe width="560" height="315" src="https://www.youtube.com/embed/OpuwY_iYctk?si=bsJA1TvVdyFInzHA" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>

比べると、mobilenetベースSSDの速さが体感できるかと思います。
懸念の精度も、アプリ側で実装している時系列処理である程度はカバーできており、動画公開用に、目立つナンバープレートを隠すだけなら、十分使えそうです。

# おわりに

今回、mobilenetベースSSDに初挑戦しました。といっても、ソフトの移植がメインで、アノテーション等はVGGベースSSDの転移学習時と全く同じものを使う等したので、mobilenet固有の学習ノウハウ等は特にありませんでした。

ただ、最初学習したときにナンバープレートが全く検知できなかったことをきっかけに、ソースを調査して、mobilenetの6箇所ある出力層の特徴マップの解像度や、DefaultBoxのサイズ定義が、小さい物体には不向きな設定になっていることを知ることができました。

DefaultBoxサイズには上述の応急措置を施しましたが、特徴マップの解像度については全く手が出せませんでした。今後は、もう少し理解を深めて、そのあたりの改良も施せたらと思いました。