PyTorchのモデルを別形式に変換する方法いろいろ（TorchScript, ONNX, TensorRT, CoreML, OpenVINO, Tensorflow, TFLite）

はじめに

本記事ではtorchvisionのresnet50を題材にPyTorchのモデルを様々な形式に変換する方法を紹介します。たくさんの種類を紹介する都合上、それぞれの細かい詰まりどころなどには触れずに基本的な流れについて記載します。また、変換後のモデルが正常に動作するかなどの確認も行いません。紹介する変換は以下の7パターンです。まさに7変化ですね！

• TorchScript
• ONNX
• TensorRT
• CoreML
• OpenVINO
• Tensorflow
• TFLite

基本環境

Ubuntu 18.04.3
PyTorch1.6

TorchScript

TorchScriptについての詳細は以下の記事でも紹介しています。詳しく知りたい方はこちらをご参照ください。

https://qiita.com/hirune924/items/8f6908e29a4e05048207
https://pytorch.org/docs/stable/jit.html?highlight=torchscript
コードは以下になります。

trans_torchscript.py

import torch
import torchvision.models as models

model = models.resnet50()
script_model = torch.jit.script(model)
script_model.save("resnet50.pt")

torch.jit.scriptのoptimize引数にダミー入力を指定するとパフォーマンスを指定の入力サイズに最適化できるようなので入力サイズが分かっている場合は指定してもいいかもしれません。
また、traceを使う場合は以下のようになります。

trans_torchscript_trace.py

import torch
import torchvision.models as models

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
script_model = torch.jit.trace(model, example_inputs=dummy_input)
script_model.save("resnet50.pt")

ちなみに以下のように読み込むことができます。

load_torchscript.py

load_model = torch.jit.load("resnet50.pt")

ONNX

DNNモデルの中間形式として一番使われているONNXについての変換をします。ONNXはversionによって対応してるオペレータが結構違うのでデプロイ先のonnxのversionに合わせてopset_versionを指定できるのは嬉しいですね。
詳細：https://pytorch.org/docs/stable/onnx.html

trans_onnx.py

import torch
import torchvision.models as models

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=True)

これをつかって推論するときはCaffe2を使うなりonnxruntimeを使うなり好きなものを使いましょう。

TensorRT

NVIDIAのGPUで動くプロダクトにデプロイするためならこのNVIDIA謹製のTensorRTを使うのがいいでしょう。
変換にはこれを使います。
https://github.com/NVIDIA-AI-IOT/torch2trt

CUDAとかTensorRTのインストールは頑張ってください。。。私はDockerを使ってnvcr.io/nvidia/tensorrt:20.07.1-py3にpytorchをインストールして使いました。
use_onnxオプションはFalseでもいいようなのですが、versionの関係で現在はTrueにしておかないと変なエラーが発生しました。
https://github.com/NVIDIA-AI-IOT/torch2trt/issues/375

trans_trt.py

import torch
from torch2trt import torch2trt
import torchvision.models as models

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
model_trt = torch2trt(model, [dummy_input], use_onnx=True)
torch.save(model_trt.state_dict(), 'resnet50_trt.pth')

ちなみに読み込むときは以下のように書けるようです。

load_trt.py

import torch
from torch2trt import TRTModule

model_trt = TRTModule()
model_trt.load_state_dict(torch.load('resnet50_trt.pth'))

CoreML

iOS端末上でモデルを動かしたいときはCoreMLを使うのがいいかもしれません。TorchScriptを使っても動かすことができますが、基本的には動かすデバイスが推奨している形式を使うのが一番いいと思います。変換にはこれを使います。
https://github.com/apple/coremltools

trans_coreml.py

import torch
import torchvision.models as models
import coremltools as ct

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
script_model = torch.jit.trace(model, example_inputs=dummy_input)
# Convert to Core ML using the Unified Conversion API
model = ct.convert(
    script_model,
    inputs=[ct.ImageType(name="input_1", shape=dummy_input.shape)])
model.save('resnet50.mlmodel')
loaded_model = ct.models.MLModel('resnet50.mlmodel')

OpenVINO

OpenVINOはIntelが提供しているライブラリです。なのでIntel製CPU上で推論を行いたいときには最も適していると思います。
OpenVINOはここからインストールしましょう。
https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit/choose-download.html

この変換はONNXを介して行います。上で作成したresnet50.onnxモデルを使用します。

trans_openvino.sh

python <INSTALL_DIR>deployment_tools/model_optimizer/mo.py --input_model resnet50.onnx

このコマンドを実行するとカレントディレクトリに色々とファイルが生成されます。ちなみに自分のInstall_dirは/opt/intel/openvino_2020.4.287/でした。

TensorFlow

deploy先がtensorflow信者なこともあるかもしれません。その場合はONNXを経由することでTensorflowの世界に行くこともできます。こちらを使います。
https://github.com/onnx/onnx-tensorflow

onnx-tf.sh

onnx-tf convert -i resnet50.onnx -o resnet50.pb

TFLite

ここまでできればtfliteの世界を垣間見ることも可能です。今回はPyTorch->ONNX->tensorflow->TFLiteという順番で変換します。ONNXを作るところからすこしポイントがあるのでそこから説明します。
まずPyTorch->ONNXですが、この時に以下のようにinput_names, output_namesを指定してやることが必要です。

convert_onnx.py

import torch
import torchvision.models as models

dummy_input = torch.randn((1, 3, 224, 224))
model = models.resnet50()
torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=True, input_names=['input'], output_names=['output'])

次にこれをTensorflowに変換します。

onnx-tf.sh

onnx-tf convert -i resnet50.onnx -o resnet50.pb

最後にこれをTFLiteに変換します。ONNX変換の時に指定したinputとoutputの名前を指定してやるのがポイントです。今回はv1の関数を使用していますが、もしv2のtf.lite.TFLiteConverterを使って変換することができた人がいたら教えてください。

convert_tflite.py

import tensorflow as tf

input_arrays = ["input"]
output_arrays = ["output"]
converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
    'resnet50.pb',
    input_arrays,
    output_arrays
)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

最後に

今回は色々な変換をサラッと試してみました。PyTorchが1.0になってDeployに力を入れ始めた時と比べると1.6になってかなり柔軟に変換できるようになってきたと思います。特にTorchScriptとONNXの恩恵ですね。基本的にC++で推論部分を記述できるならTorchScriptとLibTorchをつかえば事足りると思いますし、ONNXに対応していない推論サービスはほとんどないでしょう。PyTorchの今後の発展にますます期待が高まります！