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

coremltoolsのドキュメントを参考に、torchvisionにあるMobileNetV2をCoreMLモデルに変換してみます。

実はもうAppleが公開しているCoreMLモデルにMobileNetV2があるのですが、 これと同じモノをどのような手順で作ることができるのか を追ってみようと思います。

とにかくCoreMLモデルを作る

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

手元のMacで行う場合はminiforgeなどで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 を開くと以下のような画面が表示されます。

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

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

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

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

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

ドキュメントにもあるように、 Preview を使えるようにするには上記のフォーマットになるよう変換する必要があります。

Previewできるモデルにする

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")

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

出力の正規化

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

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

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

おわりに

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

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