はじめに

この記事では、TensorFlowモデルをONNX形式に変換し、変換されたモデルを確認して、ONNXモデルを用いた推論を行う手順について説明します。

TensorFlowモデルの準備

この章では、必要な依存関係のインストール、モデルの読み込みについて説明します。

依存関係のインストール

まずはシステムにPython 3.6以上がインストールされていることを確認してください。次に、pipを使用してTensorFlowとtf2onnxの必要なパッケージをインストールします。

$ pip install tensorflow
$ pip install tf2onnx

モデルの読み込み

TensorFlowモデルを読み込むには、tensorflow.keras.modelsモジュールのload_model関数を使用できます。この関数はSavedModel形式とHDF5形式の両方に対応しています。

SavedModelの場合は、モデルを含むディレクトリのパスを渡します。

import tensorflow as tf

model_path = "path/to/your/tensorflow/saved_model"
model = tf.keras.models.load_model(model_path)

HDF5モデルの場合は、.h5ファイルのパスを渡します。

model_path = "path/to/your/tensorflow/model.h5"
model = tf.keras.models.load_model(model_path)

モデルにカスタムレイヤーやコンポーネントが含まれる場合は、load_model関数にカスタムオブジェクト辞書を提供する必要がある場合があります。

custom_objects = {
    'CustomLayer': CustomLayer,
    'custom_metric': custom_metric
}

model = tf.keras.models.load_model(model_path, custom_objects=custom_objects)

モデルのONNX形式への変換

この章では、tf2onnxライブラリを使用してTensorFlowモデルをONNX形式に変換するプロセスと、変換プロセスのカスタマイズオプションについて説明します。

tf2onnxを使った変換

TensorFlowモデルをONNX形式に変換するには、tf2onnxモジュールのconvert.from_keras関数を使用します。この関数はTensorFlowモデルを主要な引数として受け取り、ONNXモデルを表すprotobufオブジェクトを返します。

import tf2onnx

# Perform the conversion
model_proto, _ = tf2onnx.convert.from_keras(model)

次に、protobufオブジェクトをONNXファイルとして保存します。

output_onnx_model = "path/to/output/onnx/model.onnx"

with open(output_onnx_model, "wb") as f:
    f.write(model_proto.SerializeToString())

変換プロセスのカスタマイズ

from_keras()関数に追加の引数を提供することで、変換プロセスをカスタマイズできます。よく使用される引数のいくつかは次のとおりです。

• target_opset: 使用するターゲットONNXオペレーションセットのバージョンを指定します。デフォルトはtf2onnxでサポートされる最新のオペレーションセットバージョンです。

model_proto, _ = tf2onnx.convert.from_keras(model, target_opset=12)

• large_model: モデルがメモリに収まらない場合は、このオプションを有効にしてモデルをチャンク単位で変換できます。

model_proto, _ = tf2onnx.convert.from_keras(model, large_model=True)

• output_names: ONNXモデルの出力名のリストを指定します。デフォルトでは、TensorFlowモデルからの出力名が使用されます。

output_names = ["output1", "output2"]
model_proto, _ = tf2onnx.convert.from_keras(model, output_names=output_names)

詳細なカスタマイズオプションについては、tf2onnxのドキュメントを参照してください。

https://github.com/onnx/tensorflow-onnx

変換されたONNXモデルの確認

この章では、変換されたONNXモデルを検査して検証する方法について説明します。モデルの構造を確認し、その精度をテストする方法について説明します。

モデル構造の確認

変換されたONNXモデルの構造を確認するには、onnxライブラリを使用してモデルを読み込み、一連のチェックを実行します。

import onnx

onnx_model = onnx.load(output_onnx_model)
onnx.checker.check_model(onnx_model)

これにより、モデルが無効であるか、サポートされていない操作が含まれている場合は例外が発生します。モデルがチェックをパスした場合、その精度をテストすることができます。

モデル精度のテスト

ONNXモデルの精度をテストするには、元のTensorFlowモデルと変換されたONNXモデルの両方で推論を実行し、出力を比較します。次の手順を使用できます。

1. モデルのサンプル入力を準備

import numpy as np

input_data = np.random.rand(1, 224, 224, 3).astype(np.float32)

2. TensorFlowモデルで推論を実行

tf_output = model.predict(input_data)

3. ONNX Runtimeを使用して変換されたONNXモデルで推論を実行

from onnxruntime import InferenceSession

sess = InferenceSession(output_onnx_model)
input_name = sess.get_inputs()[0].name
onnx_output = sess.run(None, {input_name: input_data})

4. 出力を比較し、差分を計算

difference = np.abs(tf_output - onnx_output).max()
print("Difference:", difference)

5. 差分が許容範囲内かどうかを確認

assert np.allclose(tf_output, onnx_output, rtol=1e-05, atol=1e-07), "Output values do not match"

変換されたONNXモデルを用いた推論

この章では、変換されたONNXモデルを使用して推論を実行する方法について説明します。ONNXモデルをONNX Runtimeで読み込み、入力データを準備し、推論を実行し、出力を処理する方法について説明します。

ONNXモデルの読み込み

変換されたONNXモデルを読み込むには、ONNX RuntimeのInferenceSessionクラスを使用します。

from onnxruntime import InferenceSession

sess = InferenceSession(output_onnx_model)

入力データの準備

モデルの入力データを準備します。入力データは、モデルの期待する入力形状とデータタイプに合わせている必要があります。例えば、次のようになります。

import numpy as np

input_data = np.random.rand(1, 224, 224, 3).astype(np.float32)

推論の実行

ONNXモデルを使用して推論を実行するには、InferenceSessionオブジェクトのrunメソッドを呼び出します。入力名を対応する入力データにマップする辞書を提供する必要があります。

input_name = sess.get_inputs()[0].name
output = sess.run(None, {input_name: input_data})

この例では、runメソッドの最初の引数としてNoneを渡しています。これは、全ての出力ノードの出力が返されることを意味します。代わりに、特定のノードの出力を取得するために、出力名のリストを提供することができます。

出力の処理

推論を実行した後、出力を処理して必要な情報や予測を取得することができます。出力の形式は、特定のモデルやタスクに依存します。例えば、モデルが分類器である場合、もっとも高い確率を持つクラスを見つけることができます。

predicted_class = np.argmax(output[0], axis=-1)
print("Predicted class:", predicted_class)

参考

https://github.com/onnx/tensorflow-onnx