技術

ONNXとは

Open Neural Network Exchange

異なる機械学習フレームワーク間でモデルを共通形式として扱えるようにする、オープンなモデル形式

ひとことで言うと

違うAIの開発ツール間でモデルをやり取りできるようにする、共通のファイル形式。

概要

ONNX(Open Neural Network Exchange)とは、PyTorch・TensorFlow等、異なる機械学習フレームワークで学習したモデルを共通の形式へ変換し、環境を問わず扱えるようにするオープンなモデルフォーマット。あるフレームワークで学習したモデルをONNX形式にエクスポートすれば、ONNX Runtime等の推論エンジンを通じて、学習時とは異なる言語・プラットフォーム・ハードウェア上でも推論できるようになる。 Microsoft・Facebook(現Meta)らが主導して策定した業界横断の標準フォーマットで、学習環境と推論環境を分離し、それぞれに最適なツールを組み合わせて使うMLOpsの実践を支える基盤の1つになっている。

背景

機械学習の開発現場では、研究・学習をPyTorchで行い、本番環境の推論では別の最適化されたランタイムを使うといった、フレームワークをまたいだ運用のニーズがある。 ONNXは、こうしたフレームワーク間の壁を取り払い、モデルを共通形式でやり取りできるよう開発された。

歴史

2017年: MicrosoftとMeta(当時Facebook)が共同でONNXを発表し、フレームワーク間の相互運用性を高めるオープンな標準として公開。以降、多くの機械学習フレームワーク・推論エンジンがONNXのエクスポート・インポートへ対応するようになった。

アーキテクチャ

ニューラルネットワークの構造(層・演算・パラメータ)を、フレームワークに依存しない共通の中間表現として定義する。モデルをONNX形式へ変換した後は、ONNX Runtime等の推論エンジンを通じて、CPU・GPU・専用アクセラレータ等さまざまなハードウェア上で実行できる。

コード例

ONNX Runtimeでの推論(概念例)

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name
output = session.run(None, {input_name: np.random.randn(1, 3, 224, 224).astype(np.float32)})
print(output[0].shape)

利点

学習で使ったフレームワークによらず、共通の形式でモデルを配布・実行できる
各種ハードウェア・言語向けに最適化された推論エンジンを通じて、実行環境の選択肢が広がる
Microsoft・Meta等が主導する業界横断の標準フォーマットで、幅広いツールが対応している

欠点

フレームワーク独自の特殊な演算やレイヤーは、ONNXへの変換で非対応になる、または挙動が変わりやすい
モデルの変換・検証に追加の作業とテストの手間がかかる
動的な制御フローを持つモデルの変換は、うまくいかない場合がある

比較

PyTorch — PyTorchで学習したモデルは、ONNX形式へエクスポートすることで他の実行環境でも利用できる
TensorRT — ONNX形式のモデルは、TensorRT等の推論エンジンへの変換元としてよく使われる

よくある質問

ONNXに変換すれば必ず高速になる？

ONNX自体は共通のモデル形式であり、高速化は組み合わせるONNX Runtime等の推論エンジンやハードウェア最適化に依存する。

TensorFlowのモデルもONNXに変換できる？

tf2onnx等の変換ツールを使うことで、TensorFlowで学習したモデルもONNX形式へ変換できる。

参考文献

Official WebsiteONNX
GitHubonnx/onnx