coremltoolsでPyTorchモデルをCore MLモデルに変換する - 2022年8月版

coremltools 4.0からはコンバーターが大きく刷新され、PyTorchモデルがONNXを経由することなく直接Core MLモデルに変換できるようになった。

coremltools 4.0でPyTorchモデルを変換する - Qiita

2022年8月現在の最新バージョンは5.2、まもなく（おそらく今年9月か10月には）6.0がリリースされる。

現行のcoremltools最新版をM1 Mac上に環境構築する手順についてはこちらに書いた。

本記事では現行の最新バージョンでPyTorchモデルをCore MLモデルに変換する実装について書く。

公式リファレンスのサンプル

まずは公式ドキュメントに記載されているサンプルコードを確認する。

Unified Conversion API

PyTorchモデルの読み込み

import coremltools as ct

# Load PyTorch model (and perform tracing)
torch_model = torchvision.models.mobilenet_v2()
torch_model.eval() 

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

Core MLモデルへの変換

# Convert using the same API. Note that we need to provide "inputs" for pytorch conversion.
model_from_torch = ct.convert(traced_model,
                              inputs=[ct.TensorType(name="input", 
                                                    shape=example_input.shape)])

Core MLモデルへの変換（ML Program）

ML Programとはなにか、については下記記事を参照：

今後はML Programがメインになっていくのでそちらがデフォルトかと思いきや、まだ旧来フォーマットがデフォルトらしい。

By default, the coremltools converter creates a neural network, but you can use the convert_to parameter to specify the mlprogram model type for an ML program model:
（デフォルトでは、coremltools コンバータはニューラルネットワークを作成しますが、convert_to パラメータを使用して ML プログラムモデル用の mlprogram モデルタイプを指定することができます。）

ML Programフォーマットのモデルに変換するには convert_to 引数に "mlprogram" を指定する。

model_from_torch = ct.convert(traced_model,
                              convert_to="mlprogram",
                              inputs=[ct.TensorType(name="input", 
                                                    shape=example_input.shape)])

シンプルな画像分類モデルを変換してみる

手元にあるPyTorchモデルとして、こちらの記事で学習・保存した画像分類（Image Classification）モデル `cifar10_0.pth` を変換してみることにした。

アーキテクチャはVGG11、CIFAR-10データセットを使用して学習したシンプルな画像分類モデル。

最小実装

公式サンプルに準じて最小限の実装を行ったコードが以下：