PytorchモデルをONNX経由でCoreMLモデルに変換する（あえてね）

CoreMLTools4.0からPytorchモデルを直接（Traced_Model経由で）変換できるようになりました。
旧式のONNX経由で変換するやりかたは非推奨になったのだけど、とはいえ「直接変換はできない」かつ「旧式ONNX方式だと変換できた」ケースもあったので、書いておくです🐣

1.PyTorchモデルをONNXにエクスポート

x = torch.randn(1, 3, 224, 224, requires_grad=True).cuda() #ダミーインプット(batch_size,color_channel,width,height) モデルがCuda使ってるなら、.cuda()

torch.onnx.export(torch_model,   # モデル
                  x,   # ダミーインプット
                  "torch_model.onnx",   # 保存パス
                  export_params=True,   # モデルの重みも保存するか
                  # opset_version=10,   # エクスポートするONNXのバージョン
                  # do_constant_folding=True,  # 最適化のために定数畳み込みを実行するかどうか
                  # input_names = ['input'],   # the model's input names
                  # output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},
                                'output' : {0 : 'batch_size'}})

2.CoreMLTools ONNX_CoreMLをインストール

pip install coremltools==3.4
pip install onnx_coreml

3.ONNXをCoreMLに変換

from onnx_coreml import convert

mlmodel = convert("torch_model.onnx")
mlmodel.save("torch_model.mlmodel")

🐣

---

お仕事のご相談こちらまで
[e-mail]

Core MLを使ったアプリを作っています。
機械学習関連の情報を発信しています。

Twitter
Medium
ホームページ