MobileNetV2をcoremltoolsでCoreMLモデルに変換する

[coremltoolsのドキュメント](https://apple.github.io/coremltools/docs-guides/source/convert-a-torchvision-model-from-pytorch.html)を参考に、torchvisionにあるMobileNetV2をCoreMLモデルに変換してみます。

実はもう[Appleが公開しているCoreMLモデル](https://developer.apple.com/jp/machine-learning/models/)にMobileNetV2があるのですが、 **これと同じモノをどのような手順で作ることができるのか** を追ってみようと思います。

## とにかくCoreMLモデルを作る

Google Colaboratoryでも手元のMacでもよいので、以下を実行します。

> 手元のMacで行う場合はminiforgeなどでPython環境を作るのがよさそうです

```python
import torch
import torchvision

# 学習済みMobileNetV2モデルをロード
torch_model = torchvision.models.mobilenet_v2(pretrained=True)

# 推論モードに設定
# https://stackoverflow.com/questions/60018578/what-does-model-eval-do-in-pytorch
torch_model.eval()

# モデルをトレースする
# 入力テンソルはMobileNetV2の仕様に準拠したランダムな値
example_input = torch.rand(1, 3, 224, 224) 
traced_model = torch.jit.trace(torch_model, example_input)
out = traced_model(example_input)

# CoreMLモデルに変換
model = ct.convert(
    traced_model,
    convert_to="mlprogram",
    inputs=[ct.TensorType(shape=example_input.shape)]
)

# CoreMLモデルを保存
model.save("MobileNetV2.mlpackage")
```

できあがった `.mlpackage` を開くと以下のような画面が表示されます。

![スクリーンショット 2024-01-08 15.45.32.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25727/8bf3ffa3-8666-239b-7ade-aed5ef96d068.png)

中身はよくわかりませんが、モデルの変換はできていそうです。

## Appleが公開しているモデルと比較する

Appleが公開している `MobileNetV2.mlmodel` を開くと `Preview` というタブがあり、ここに画像を放り込むことでモデル動作検証が行えるようになっています。

![スクリーンショット 2024-01-08 15.57.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25727/72d3ab69-1716-c35e-10f5-7d14e4089637.png)

一方自分で作成した `MobileNetV2.mlpackage` を開くと `Preview` タブは表示されません(前項のスクリーンショット参照)。
モデルの概要と `Predictions` タブの中身を比較すると、以下の違いがあることがわかります。

* `Model Type` が `Neural Network Classifier` でない
* `Input` が `Image(Color 224 x 224)` でない
* `Output` が `Dictionary (String → Double)` でない

[ドキュメントにもあるように](https://apple.github.io/coremltools/docs-guides/source/xcode-model-preview-types.html)、 `Preview` を使えるようにするには上記のフォーマットになるよう変換する必要があります。

## Previewできるモデルにする

```python
import torch
import torchvision

torch_model = torchvision.models.mobilenet_v2(pretrained=True)

torch_model.eval()

example_input = torch.rand(1, 3, 224, 224) 
traced_model = torch.jit.trace(torch_model, example_input)
out = traced_model(example_input)

# 分類ラベルをダウンロード
# https://apple.github.io/coremltools/docs-guides/source/convert-a-torchvision-model-from-pytorch.html#download-the-class-labels
import urllib
label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
class_labels = urllib.request.urlopen(label_url).read().decode("utf-8").splitlines()
class_labels = class_labels[1:] # remove the first class which is background
assert len(class_labels) == 1000

# 入力をテンソルからImageに変更
# 合わせてtorchvisionモデル用の前処理を追加
# https://apple.github.io/coremltools/docs-guides/source/convert-a-torchvision-model-from-pytorch.html#preprocess-the-image-input-for-torchvision-models
image_input = ct.ImageType(
  name="input_1",
  shape=example_input.shape,
  scale=1/(0.226*255.0),
  bias=[-0.485/(0.229), -0.456/(0.224), -0.406/(0.225)]
)

model = ct.convert(
    traced_model,
    # Model Type: Neural Network Classifier
    convert_to="neuralnetwork",
    # Input: Image(Color 224 x 224)
    inputs=[image_input],
    # Output: Dictionary(String → Double)
  	classifier_config=ct.ClassifierConfig(class_labels)
)

model.save("MobileNetV2.mlpackage")
```

![スクリーンショット 2024-01-08 15.59.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25727/1c597ea2-7d74-e02f-a9bf-71116b618094.png)

今度は `Preview` タブが表示されるようになりました！
試しに画像を放り込んでみると、分類がうまくいっていないだけでなく、なぜか100%を超えちゃっています🤔

## 出力の正規化

これは **torchvisionモデルが、出力の正規化を行っていないことが原因** でした。
https://pytorch.org/hub/pytorch_vision_mobilenet_v2/
なので、以下のようにモデル出力の後に[Softmax](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html)関数を追加し、正規化された出力を得られるようにします。

```python
torch_model = torch.nn.Sequential(
    torch_model,
    torch.nn.Softmax(dim=1)
)
```

これで変換したCoreMLモデルで試してみると、今度は100%を超えないようになりました！
結果や精度もAppleのものに近いので問題なさそうです。

![スクリーンショット 2024-01-08 15.56.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/25727/06edd0a1-2886-54a7-a9ec-f6dd2161ea47.png)

## おわりに

とりあえず既存モデルをcoremltoolsを使ってCoreMLモデルに変換できることと、その際の注意点を理解できました。これで転移学習やファインチューニングなりで作ったカスタムモデルでも、同じワークフローでCoreMLモデルを作り出すことができそうです。

また別の記事で、カスタムモデルとVision Frameworkを使った事例を紹介しようと思います。