🔦

coremltoolsでPyTorchモデルをCore MLモデルに変換する

2023/12/15に公開

iOS

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

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

https://note.com/shu223/n/n95af3b29d89e

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

公式ドキュメントの変換サンプル

まずは公式ドキュメントに記載されているサンプルコードを読み、変換の実装を確認する。

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 を変換してみることにした。

https://note.com/shu223/n/n1fd46238a0c9

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

最小実装

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

import os
import torch
from model import VGG11
import coremltools as ct

model_name = "cifar10_0"

# PyTorchモデルの読み込み
model = VGG11()
model_state_dict_path = os.path.join(model_dir, model_name + '.pth')
model.load_state_dict(torch.load(model_state_dict_path))
model.eval()

example_input = torch.rand(1, 3, 32, 32)
traced_model = torch.jit.trace(model, example_input)

# Core MLモデルへの変換
mlmodel = ct.convert(traced_model,
                     convert_to="mlprogram",
                     inputs=[ct.TensorType(name="input", 
                                           shape=example_input.shape)])

# 保存
mlmodel.save(model_name + '.mlpackage')

生成されたモデルをXcodeでプレビューしてみる：

Generalタブ

Predictionsタブ

問題点

Xcodeプレビューで、以下の問題があることがわかる。

入力も出力もMultiArray (MLMultiArray)
画像分類モデルとして認識されていない
Previewタブがない → Xcode上でモデルを試せない

変換オプションの指定

問題点を解消していくため、Core MLモデルの変換オプションを指定する。ドキュメントはこちら。

入力を画像にする

ドキュメント：

coremltools 6の場合、convertメソッドのinputs引数と、colorlayout を利用してこう書けるらしい。

mlmodel = ct.convert(traced_model,
                     convert_to="mlprogram",
                     inputs=[ct.ImageType(shape=example_input.shape,
                                          color_layout=ct.colorlayout.RGB,)])

ただ今はcoremltools 5.xを使用しているので、エラー。

AttributeError: module 'coremltools' has no attribute 'colorlayout'

color_layout の指定をシンプルに省略してみる。

mlmodel = ct.convert(traced_model,
                     convert_to="mlprogram",
                     inputs=[ct.ImageType(name="input",
                                          shape=example_input.shape)])
mlmodel.save(model_name + '.mlpackage')

Xcodeでプレビューすると、

入力がImageになった🎉

ただ、まだ出力はMultiArrayであり、Xcodeは本モデルを画像分類モデルとして扱っていない。Previewタブもない。

クラスラベルを追加する

ドキュメント：

以下のようにconvertメソッドに classifier_config を指定。

class_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
classifier_config = ct.ClassifierConfig(class_labels)

mlmodel = ct.convert(traced_model,
                     convert_to="mlprogram",
                     inputs=[ct.ImageType(name="input",
                                          shape=example_input.shape)],
                     classifier_config=classifier_config)


mlmodel.save(model_name + '.mlpackage')

Xcodeでプレビューしてみると、

Class Labelsが入った🎉

しかしまだPreviewタブが現れない

ML Programフォーマットをやめてみる。

mlmodel = ct.convert(traced_model,
                     inputs=[ct.ImageType(name="input",
                                          shape=example_input.shape)],
                     classifier_config=classifier_config)

mlmodel.save(model_name + '.mlmodel')

Xcodeのプレビュー上でPreviewタブが現れ、画像を入力して試せるようになった！

ML ProgramとXcode Preview

使用しているXcodeのバージョンは13.4.1。

ML Programは2021年に登場したばかりなので、Xcode 13.x ではまだML Programモデルのプレビューをサポートしてないのかもしれない。

・・・と思いXcode 14.0 beta 2でプレビューしてみたところ、やはりPreviewタブは現れなかった（.mlmodelだと現れた）。

model.preview.type を設定する必要があるのかもしれない。

Classifierモデルでは設定しないでいいと明記されているのだが...

Some model architecture types, such as Neural Network Classifier, don't require a model.preview.type, and some model preview types don't require preview parameters.

補足： PyTorchモデルの種類

この部分はcoremltoolsは関係ないのだが、PyTorchのモデル .pth ファイルには2種類ある。

state_dict を保存したものと、モデル全体を保存したものだ。

state_dict とは、『各レイヤーをそのパラメーターテンソルにマップする単純なPythonディクショナリ』らしい。

こちらのタイプのモデルファイルは、モデル定義クラスから初期化して、load_state_dict メソッドを用いて state_dict をファイルから読み込む。

model = MyModelDefinition(args)
model.load_state_dict(torch.load('load/from/path/model.pth'))

一方のモデル全体を保存するタイプの方は、torch.loadするだけだ。

model = torch.load('load/from/path/model.pt')

Python / PyTorch素人からすると後者の方が楽で良さそうだが、PyTorchオフィシャルの推奨としては前者らしい。

理由としては以下のように、もともとのパス構造に依存してしまい、柔軟性が悪くなるらしい。

Pythonのpickleモジュールは内部で使用されるため、シリアル化されたデータは特定のクラスにバインドされ、モデルの保存時に正確なディレクトリ構造が使用されます。 Pickleは、特定のクラスを含むファイルへのパスを保存するだけです。これは、読み込み時に使用されます。

ご想像のとおり、保存されたモデルが同じパスにリンクしていない可能性があるため、リファクタリング後にコードが破損する可能性があります。パス構造を維持する必要があるため、このようなモデルを別のプロジェクトで使用することも困難です。

参考：

iOSアドベントカレンダー

本記事はiOS Advent Calendar 2023 シリーズ2の11日目 ^[1]の記事です。

脚注

12/15時点で空いていたので書きました。 ↩︎