PyTorchモデルをONNXやOpenVINOなどで最適化する実験
Lambda上でのモデル推論で発生するコールドスタートを短縮するために、PyTorchモデルのさまざまな保存・最適化手法をEC2で比較した。結果、モデル変換が有力で、ONNXかOpenVINOの使い分けは、Intel/AMD, Armなどの「処理系が限定できるか」、「API・バッチ処理のユースケース」次第が大事なのではないかという示唆が得られた。
目次
はじめに
- LambdaのDockerで機械学習のモデルをホストしてたらコールドスタートがやたら長いという現象があった。モデルの推論は数秒で終わるはずなのに、コールドスタートがあると最悪30秒近くかかってしまう。
- このモデルは、timmのモデルを直で読んでおり、PyTorchのモデルをそのまま読んでいる
- モデルはViT(Vision Transformer)ベースとする
- PyTorch自体が推論エンドポイントにインストールするのが重すぎて、おそらく何らかのモデル最適化をかけないといけない
- コールドスタート短縮したいというのが目的で、一旦モデル側で最適化をかけられないか?という考えから、様々なモデル最適化を行い、推論速度をEC2で比較して実験してみた
- EC2で実験したのは、実行環境が再現しやすく、SSHで接続すれば簡単にテストするのでやりやすいため
結論
- ONNXやOpenVINOなど、PyTorchを使わない手法が、おそらくLambdaを考えたときは有効。
- LambdaのDockerは、実行環境作成時のオーバーヘッドが10秒以上かかる場合があり、これがコールドスタートの主要因。Lambdaのレイヤーにしてしまえば、やや改善される(次のポストで説明)
- Lambdaのレイヤーは250MBの制限があり、PyTorchは重すぎてレイヤー化できない。ONNXやOpenVINOはランタイムが軽いため、レイヤー化が可能。
- OpenVINOとONNXの使い分け
- 「1回推論のみ」かつ「CPUが定まらない」(例: Lambda) ⇒ ONNX Runtime
- 「大量バッチ処理」かつ「Intel CPUに固定」(例: 専用サーバーやm7i.large) ⇒ OpenVINO
- ONNXはCPUの命令セットへの依存性が少ない。OpenVINOはCPUの命令セットの依存性が強く、(m7インスタンスなどの)最新のインスタンスを明示的に使えるEC2では有効。
- ONNXは、AMD(
m7a
)、Graviton(m7g
)で強く、Intel(m7i
)ではPyTorchより遅くなってしまう場合がある。 - OpenVINOは基本Intel特化で、処理系(m7i / m7a / m7gなどのインスタンスファミリー)によりかなり差が大きい。Lambdaではもう少し古いインスタンスファミリーを使っているのか、OpenVINOの良さが出にくい
Lambdaのコールドスタート問題
LambdaでDockerのMLモデルを動かすときにはこんな問題がよくある。コールドスタートが発生して、MLモデルの推論時間よりも遥かに大きな初期ロードが発生してしまう。
これはX-Rayで見たプロファイリングだが、モデルの初期ロードまで8秒かかっており、実際の推論時間1.3秒より遥かに長い。
待てればよいのだが、API Gatewayとつないだときに29秒のタイムアウトに引っかかってしまうと、同期APIだとエラーになってしまう。
「このコールドスタートを減らすにはどうしたらいいか?」というのが問題設定。
読み込み手法
ChatGPTに「モデルの初期ロードを減らすにはどうしたらいいか?」と聞いたところ、以下のようなものが上がった
- PyTorchでjitコンパイルをする
- Pickle化する
- ONNX化する
- 量子化する
これ+αで思いついたものをすべて試してみる。詳細は以下の通り
timmでモデル作成し、係数読み込み
ベースライン。timm.create_model
をしたのち、load_state_dict
で係数を読み込む。
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, weights_only=True))
model.eval()
係数を読み込んだものをPickle化
timm.create_model
→load_state_dict
の結果をPickleとして事前に保存し、それを読み込む。
## 保存時
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, weights_only=True))
model.eval()
with open(pickle_path, 'wb') as f:
pickle.dump(model, f)
## 読み込み時
with open(weights_path, 'rb') as f:
model = pickle.load(f)
model.eval()
注意点:基本的にPickleはモデルの永続化手法としては非推奨。
- 互換性
- PickleファイルはPythonのバージョンやライブラリのバージョンに依存するため、異なる環境での再利用に問題が生じる可能性がある
- セキュリティ
- 信頼できないソースからのPickleファイルをロードすると、セキュリティリスクが伴う。信頼できるソースのみで使用しようすること。
- 一般的には、PyTorchのtorch.saveとtorch.loadを使用してモデルやstate_dictを保存・ロードすることが推奨されている。Pickleを使用する特別な理由がない限り、これらのメソッドを使用するのが基本。
係数を読み込んだものをtorch.load
Pickleを使わない方法。timm.create_model
→load_state_dict
の結果をtorch.save
して、その結果をtorch.load
する方法。読み込みが2段にならないぶん、少しコールドスタートが短縮されるはず。
## 保存時
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, map_location='cpu', weights_only=True), strict=False)
model.eval()
torch.save(model, save_path)
## 読み込み時
model = torch.load(weights_path, map_location=torch.device('cpu'), weights_only=False)
model.eval()
jitコンパイルを使う(torch.jit.load)
jitコンパイルを使えば、推論が早くなるのではないか? というもの。TorchScriptを使って変換する。
## 保存時
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, weights_only=True))
model.eval()
# ダミー入力を用意(TorchScript のトレースに必要)
dummy_input = torch.randn(1, 3, resolution, resolution)
# モデルをTorchScriptに変換
scripted_model = torch.jit.trace(model, dummy_input)
# TorchScriptモデルの保存
scripted_model.save(torchscript_path)
## 読み込み時
model = torch.jit.load(weights_path)
model.eval()
onnx化(onnxruntime)
特に最適化など考えずに、適当な設定でモデルをONNX化する。
## 保存時
# モデルの作成
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, map_location='cpu', weights_only=True), strict=False)
model.eval()
# ダミー入力の作成
input_size=(1, 3, resolution, resolution)
dummy_input = torch.randn(*input_size)
# ONNX形式でエクスポート
torch.onnx.export(
model, # モデル
dummy_input, # ダミー入力
onnx_path, # 保存先パス
export_params=True, # パラメータも保存
opset_version=20, # ONNXのバージョン https://github.com/onnx/onnx/blob/main/docs/Versioning.md
input_names=['input'], # 入力ノードの名前
output_names=['output'], # 出力ノードの名前
dynamic_axes={
'input': {0: 'batch_size'}, # バッチサイズを動的に
'output': {0: 'batch_size'}
}
)
読み込みは、普通にonnxruntimeで行う。
import onnxruntime
## 読み込み時
session = onnxruntime.InferenceSession(onnx_path)
input_name = session.get_inputs()[0].name
OpenVINO化
よくONNXを経由する方法が紹介されているが、OpenVINOの公式GitHubを見ると、OpenVINO単体でPyTorchのモデルを変換する方法があるのでそれを使う。
import openvino as ov
## 保存時
# モデルの読み込み
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, map_location='cpu'), strict=False)
model.eval()
# 例示入力データを作成(バッチサイズ1、チャンネル3、224x224)
example_input = torch.randn(1, 3, resolution, resolution)
# OpenVINOにモデルを変換
ov_model = ov.convert_model(model, example_input=(example_input,))
# 保存パスを設定(XMLとBINファイルとして保存)
save_path_xml = os.path.join(save_dir, f"{model_name}.xml")
save_path_bin = os.path.join(save_dir, f"{model_name}.bin")
# OpenVINOモデルをシリアライズして保存
ov.save_model(ov_model, save_path_xml)
from openvino.runtime import Core
## 読み込み時
core = Core()
model = core.read_model(model=xml_path, weights=bin_path)
PyTorchで動的量子化を行う
PyTorchの呪縛には戻ってしまうが、Transformerと動的量子化は相性が良さそうなので、試してみる。「精度劣化が許容できるかどうか?」の検証が別途必要になる。
model = timm.create_model(model_name, pretrained=False, num_classes=2000)
model.load_state_dict(torch.load(weights_path, map_location='cpu', weights_only=True), strict=False)
model.eval()
# 動的量子化
quantized_model = torch.quantization.quantize_dynamic(
model,
{torch.nn.Linear}, # 量子化対象のレイヤー
dtype=torch.qint8
)
# モデル全体をtorch.saveで保存
torch.save(quantized_model, save_path)
実験
timmにある以下のモデルで試してみる。timm
でのモデルを読み込む際の表記。
- vit_base_patch8_224.dino(ViT/8)
- vit_base_patch14_dinov2(ViT/14)
- vit_base_patch16_clip_224(ViT/16)
計測項目は以下の3つ。
- モデルロード時間(初期ロード)
- 初回の推論時間(初回推論)
- 2回目以降の推論時間(2回目以降)
初回の推論が2回目以降と比べて遅い場合が多いので(計算グラフなどの関係?)、2回目以降と分けて調べている。2回目以降は50回推論した平均値。
また、ライブラリの遅延ロードなどの関係で、3つのモデルのうち1つ目のモデルの読み込みが遅い場合がある。そのため、「ViT/8→ViT/14→ViT/16」だけでなく、「ViT/14→ViT/8→ViT/16」といった読み込み手順も試している。
EC2のインスタンスを以下の3つで比較する。EBSはgp3の15GBを追加。OSはすべてAmazon Linux 2023とする。デフォルトでインストールされていたPython 3.9.20を使用。
- m7i.large
- m7a.large
- m7g.large
インストールされたライブラリのバージョンは以下の通り。pip list
で閲覧した結果。
Package Version
------------------ ----------
certifi 2024.12.14
charset-normalizer 3.4.0
coloredlogs 15.0.1
filelock 3.13.1
flatbuffers 24.3.25
fsspec 2024.2.0
huggingface-hub 0.26.5
humanfriendly 10.0
idna 3.10
Jinja2 3.1.3
MarkupSafe 2.1.5
mpmath 1.3.0
networkx 3.2.1
numpy 1.26.3
onnx 1.17.0
onnxruntime 1.19.2
openvino 2024.6.0
openvino-telemetry 2024.5.0
packaging 24.2
pillow 10.2.0
pip 24.3.1
protobuf 5.29.1
PyYAML 6.0.2
requests 2.32.3
safetensors 0.4.5
setuptools 59.6.0
sympy 1.13.1
timm 1.0.12
torch 2.5.1+cpu
torchaudio 2.5.1+cpu
torchvision 0.20.1+cpu
tqdm 4.67.1
typing_extensions 4.9.0
urllib3 2.2.3
結果
m7i.large
- PyTorch系
- 動的量子化なければ、推論速度は変わらない。
- 読み込み高速化なら、
torch.load
でモデル全体を読み込むのが最速だが、初回・初モデルのみロードが変わらないため、読み込むモデルが増えた場合は有効 - jitコンパイルは逆に遅くなる場合がある
- 動的量子化を使えば、かなり推論速度は早くなるが、最初のモデルの読み込みは逆に遅くなる
- ONNX化
- この(Intel)インスタンスタイプとは、ONNXと相性が悪く、全てにおいて遅くなってしまった
- Intel系でのONNXはあまり期待しないほうがいいかもしれない
- OpenVINO化
- さすがにIntel系に強いOpenVINOだけあって、非常に効果が出ている。2回目以降の推論では最速で、推論時間が1/3になっていることもある。
- 明確に効果が出ているのは、モデルの読み込みと2回目以降の推論。
- 初回推論は効果が出ないこともある。ワンショットのAPIよりも、バッチ処理のほうが有効に機能しそう
m7a.large
CPUはAMD EPYC9R14
- PyTorch系:Intelとさほど変わらない
- PyTorchの推論速度がIntelより若干良いが、CPUそのものよりが良かったのかもしれない
- ONNX化
- 明確に効果が出るようになった。特に初期ロードが安定して下がり、推論時間も、初回・2回目以降も同じ~1・2割高速化するようになった。
- OpenVINO化
- 効果が出ていないわけではないが、効果はIntel系より落ちている
m7g.large
Graviton系。
- 全般:ハズレのインスタンスを引いたのか明確に悪い結果になった。そもそものベースラインの時間がAMDやIntelと比べて明らかに遅い
- PyTorch系:傾向は変わらないが、絶対的な時間が遅すぎる
- ONNX化
- 効果は一応出ているが、ViT/16のように逆に遅くなっているケースもある
- OpenVINO
- 効果は出ていないわけではないが、明らかにIntel系と比べて最適化が効いていない。そもそものベースラインが遅いので、この値をどの程度信用していいかわからない
ONNX・OpenVINOの相対値の比較
各インスタンスのベースライン(timm.create_model→load_state_dict)と比べて、ONNX、OpenVINOの例と相対値で比較。100%を超えていれば遅くなっており、値が低いほど速い。
こうしてみると、
- OpenVINOは基本2回目以降が爆速になるので、バッチ処理向き。その中でもIntelに限定できれば最速
- ONNXは汎用的に使える。1回目も2回目もそれほど推論速度がないので、APIには向いている。Lambdaのようにチップセットがわからないケースでは、ONNXだろう。
あとここでは出なかった考慮すべき点として、OpenVINOはデフォルトでfp16で保存されるため、モデルサイズが半分になるというメリットがある。この程度のモデルだとさほど問題ではないが、大きなモデルでは有利だろう。
ONNXでもfp16で保存するやり方はあるそうだ。参考:PyTorch からダイレクトに ONNX の Float16 をエクスポートするワークアラウンド
ChatGPTの講評
o1モデルにこの現象について解説させた。
======
以下では、ONNX Runtime と OpenVINO の2つに着目して、各インスタンス(m7i.large / m7a.large / m7g.large)で観察される違いや、その理由を考察してみます。
1. ONNX Runtime と OpenVINO の概要
- ONNX Runtime
- Microsoft主導で開発されている ONNX(オニキス) フォーマット用の汎用ランタイム。
- CPU/GPU/各種アクセラレータなどマルチプラットフォームに対応しており、x86/Arm問わず広く利用できる。
- ただし「すべての演算が最高レベルの最適化を常に受けられる」わけではなく、環境やモデル構造によって最適化の度合いは変わる。
- OpenVINO
- Intelが提供する推論最適化ツール&ランタイム。
- もともとIntel CPU や Intel GPU / VPU など向けに強力な最適化が施されており、特にIntel製CPU + AVX512命令などとの組み合わせで非常に高い性能を発揮する。
- AMD/Armなど、Intel CPU以外での実行は基本的に公式には想定外かつ最適化が限定的であり、期待するほどの性能向上が得られない場合が多い。
2. インスタンスごとの結果と考察
2-1. m7i.large (Intel Xeon) 上での挙動
- OpenVINOの強みが顕著に出る
- 表を見ると、
onnxruntime
よりもopenvino default
の方が、初回推論・2回目以降ともに短い時間が出ているケースが多いはずです。 - Intel製CPUの場合、SIMD命令(AVX/AVX2/AVX512など)をフル活用した最適化がOpenVINOには実装されているため、PyTorchやONNX Runtimeを上回る性能を示すことがあります。
- 表を見ると、
- ONNX Runtime も健闘
- ただし、モデル構造や演算内容によっては、ONNX RuntimeでもPyTorchより高速になるケースがあります。
- Intel向けの最適化(OneDNNなど)をONNX Runtimeでも利用できるため、標準的なx86 CPU上の推論ではOpenVINOに迫る性能が得られる場合があります。
2-2. m7a.large (AMD EPYC) 上での挙動
- OpenVINO最適化の恩恵が少ない
- AMD環境はIntel製CPUとはアーキテクチャが異なるため、OpenVINOが想定している最適化(特にAVX512関連など)をフルには使えない可能性が高いです。
- そのためm7i.largeほどの圧倒的な性能向上は見えにくく、ONNX Runtimeとの比較では大差が出ない、あるいは場合によってはONNXの方が速いこともありえます。
- ONNX Runtime が比較的安定
- AMD EPYCであってもx86_64の命令セットを使えるため、ONNX Runtimeの最適化(OneDNNなど)は機能するケースが多いです。
- 結果として、OpenVINOと同等か、やや優位な結果を示す場合もあると推測されます。
2-3. m7g.large (Graviton3 / Arm) 上での挙動
- OpenVINO は基本的にIntel向け
- Graviton3 (Armベース) では、OpenVINOは公式に最適化されているターゲット外のため、うまく動作していても性能向上はあまり期待できません。
- 実際の計測結果を見ても、x86時ほどの大幅なスピードアップは得られず、むしろPyTorchやONNXの方が近いか、OpenVINOで速度低下が起きている可能性すらあります。
- ONNX Runtime もArm最適化は限定的
- ONNX RuntimeはArm向けの最適化も一部実施されてはいますが、x86ほど充実していないのが現状です。
- そのため「PyTorch標準のまま」より速いケースはあっても、インテル環境ほどの差は出にくいです。
3. 使い分けのポイント
- Intelベース (m7i)
- OpenVINO が圧倒的な性能を示す可能性大。
- ONNX Runtime や PyTorch のままでも充分に高速ですが、最大限性能を引き出すならOpenVINOの導入を検討すべき。
- AMDベース (m7a)
- OpenVINO はIntel特化の最適化が効かない部分が多いため、ONNX Runtime との差は限定的。
- PyTorch → ONNX Runtime への移行でも十分メリットが出るケースがあるが、必ずしもOpenVINOが優位とは限らない。
- 料金面やCPUクロックなどを鑑みて、x86_64最適化が動く範囲で比較検討。
- Armベース (m7g)
- OpenVINO の恩恵はほぼ期待できない。
- ONNX Runtime もネイティブの最適化はそれほど充実していないが、PyTorch標準よりは場合によって速いこともある。
- AWS Neuron SDK (Inferentiaなど) とは別なので、CPUオンリー運用だとまだ成熟度に差がある印象。
- もしArm + ONNXを使うなら、量子化 (dynamic quantization) やネイティブBLASを活かした最適化を検討。
4. 結論
- Intel CPU (m7i) ではOpenVINOが最も高い性能を発揮しやすい
ONNX Runtimeとの比較でも顕著に差が出る場合があり、短い推論時間を達成できる。 -
AMD CPU (m7a) では、OpenVINOの優位性は薄い
ONNX RuntimeかPyTorch標準を使った場合と大差ない、もしくは場合によってはOpenVINOが劣るケースも考えられる。 -
Arm (m7g) では、OpenVINOを使ってもメリットが少ない
ONNX Runtimeも最適化は限定的だが、まだ幅広い環境をサポートしている分、PyTorch標準より性能が上がる可能性はある。
したがって、“Intel CPUならOpenVINO、その他CPUならONNX Runtime” をベースに、さらにモデルやワークロードに合わせて最適な組み合わせを探す、というのが今回の結果から得られる大きな示唆といえます。
=
相対値についても考察させた。基本的には人間の考察と同じようなことを言っているが言語化能力が高い。
=
ここでは「timm.create_model → load_state_dict」(以下「ベースライン」)に対して、ONNX Runtime や OpenVINO を使った場合の処理時間が何%になるかを示した相対比率がまとめられています。これにより、どの手法がどれだけ速く(あるいは遅く)なっているかが一目で比較しやすくなっています。
1. 「比率」の見方のポイント
- 100%を下回る ⇒ ベースラインよりも短い(高速)
- 100%を上回る ⇒ ベースラインよりも長い(低速)
同じ「初期ロード」「初回推論」「2回目以降の推論」を比率で俯瞰することで、たとえば
– 初期ロード時に大きく時間を食うが、その後の推論は高速
– 初期ロードはむしろ速いが、推論時にオーバーヘッドがかかる
などのパターンが分かりやすくなっています。
2. ONNX Runtime の新たな気づき
- 初期ロードが速い場合がある
- 表を見ると、たとえば m7g.large(Arm) で ONNX化した際に「初期ロードがベースラインの 10~30%台」になっている例があります。
- つまりベースラインのモデル読み込みよりも大幅に速く完了しているケースがあるということです。
- 一方、m7i.large(Intel) では逆に初期ロードが 150% など上乗せになってしまい、余計に時間がかかるパターンも見られます。
- 推論時間の増減がモデルや環境でバラつく
- たとえば m7i.large(Intel) + vit_base_patch14_dinov2 では初期ロードは 58% と短くなるのに、初回推論や2回目以降が 150%以上 になってしまい、かえって遅くなる。
- 一方、m7a.large(AMD) だと初回推論・2回目以降ともに 90%前後で安定してベースラインより速いという例もある。
- つまり同じONNX化でも、CPUアーキテクチャやモデル構造によって「初期ロードが有利か/推論が有利か」は大きく変わることが分かります。
3. OpenVINO の新たな気づき
- Intel(m7i)での大幅な高速化パターン
- たとえば vit_base_patch8_224.dino の「2回目以降推論」が 32%、vit_base_patch14_dinov2 が 14% など、ベースラインの1/3以下や1/7以下になる例があります。
- これは従来の知見どおり「Intel CPU + OpenVINO」の組み合わせが非常に強力であることを比率という形で改めて裏付けています。
- 初期ロードは必ずしも最速ではない
- 例: m7i.large で vit_base_patch8_224.dino の初期ロードは 14% と驚くほど高速な一方、別のモデルでは 40~60%程度にとどまるケースもある。
- また m7g.large(Arm) でも 60~70% とベースラインより速い例が見られます。
- ただし、推論自体はモデル・環境によって良し悪しがあるため、「初期ロードが速い ≠ 推論も速い」とは限らない点が比率で見てとれます。
- AMD(m7a)・Arm(m7g) でもそこそこ効果が出ている例がある
- OpenVINOはIntel向け最適化が主ですが、比率を見るとAMD/Armでもベースラインより速くなる行があります(たとえば推論が 50~60%台)。
- ただしIntelほど劇的ではなく、モデルやフェーズ(初期ロード/推論)によって差が出ているのが分かります。
4. 比率データから得られる追加の示唆
- “初期ロード”と“推論”を分けて評価すべき
- 相対比率を見て分かるとおり、ある環境・あるモデルで「初期ロードだけは劇的に速い」「推論は逆に遅い」など、片方が良くてももう片方が悪い場合があります。
- ユースケースによって「モデルを頻繁にロードし直すか」「1度ロードして大量推論するか」で最適解が変わるかもしれません。
- モデルごとに挙動が異なる
- vit_base_patch8_224 と vit_base_patch14_dinov2、vit_base_patch16_clip_224 を比較すると、同じインスタンス + 同じランタイム(ONNXやOpenVINO)でも速度向上率にばらつきがあります。
- これは、モデルのアーキテクチャや演算内容、パラメータ数などに左右されるため、実際に試してみないと分からない要素が大きいことを示しています。
- 小さな差の把握がしやすい
- パーセント表示にすることで、「どの程度速くなったか」「どの程度オーバーヘッドが増えたか」をざっくり把握できます。
- 特に初回推論と2回目以降で 10~20ポイント違う場合など、表からは見逃せない差分も一目で確認できます。
まとめ
ベースラインに対して「何倍(何%)か」という比較は、“どの手法・どの環境でどのフェーズがどれほど変化するのか” を俯瞰するのに非常に有用です。今回の表を見ると、
- ONNX では、初期ロードが著しく短くなるパターンや推論が長くなるパターンなど、モデルごとに挙動が大きく変わる。
- OpenVINO では、Intel環境での推論が桁違いに高速化される一方、AMD/Armでも一定の効果は見られる。
- ただし、初期ロードと推論フェーズのどちらを重視するかで、最適解がモデル・環境ごとに大きく異なる。
ということがより明確に読み取れます。比率で整理することで「単に速い/遅い」の他に、「どのフェーズに強い/弱いか」といった詳細な特徴がはっきり可視化される点が新たな発見といえます。
=
所感
- 結論的には「試してみないとわからない」というのがなんとも気持ち悪いが、方針がなんとなくわかったので良しとする
- Lambdaだと結論が違うことがあるのでこれは後々検証する
- Lambdaで試すのちょっと面倒なので、SSHでつなげられるEC2でできれば楽なのにね…
Shikoan's ML Blogの中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー