機械学習フレームワークを使用する際、学習環境と推論環境が異なる場合があります。たとえば、画像認識モデルをCaffeで学習したものの、本番環境ではTensorFlowで推論を行うといったケースです。

このような状況に対応するために、異なる学習フレームワークで作成されたモデル間の相互変換が求められます。モデル変換には主に「直接変換」と「規格化変換」の2つのアプローチがあります。この記事ではそれぞれの変換方法と技術的な詳細について詳しく説明します。

直接変換

直接変換は、AIフレームワークから生成されたモデルを、ターゲットフレームワークで使用可能な形式へと変換する手法です。例えば、MindSpore ConverterはMindSporeのモデルを推論エンジンのIR形式に直接変換します。

規格化変換

規格化変換では、オープンなファイル規格を定義し、主要なAIフレームワークがその規格をサポートできるように設計されています。たとえば、PyTorchのモデルを直接変換するのではなく、PyTorchをONNX形式に変換し、その後ONNX Converterを使って推論エンジンのIR形式に変換します。これにより、多くの主要なAIフレームワークがこの仕組みに対応しています。

直接変換の手順

1. データ読み込み：AIフレームワークによって生成されたモデルファイルを読み込み、テンソルの型・形式、演算子の種類とパラメータ、計算グラフ構造と命名規則、およびそれらの関連情報を識別します。
2. 形式変換：第1段階で取得したモデル構造とパラメータ情報を、推論エンジンが対応している形式へとコードレベルで翻訳します。複雑な演算子については、Converter内部で変換関数を定義して対応します。
3. モデル保存：推論エンジン上でモデルを保存し、推論エンジンで利用可能な形式のモデルファイルを生成します。

変換プロセスでは、各フレームワークにおける演算子実装の違いや、テンソルデータの格納形式（NCHW、NHWCなど）の互換性、演算子のパラメータ名や意味の違い、および変換後のモデルのパフォーマンス向上のための最適化（演算子融合、定数畳み込みなど）などの課題を考慮する必要があります。

直接変換の実例

以下のコードは、事前学習済みのTensorFlowモデルを読み込み、それをPyTorchモデルへと直接変換する例です：

import tensorflow as tf
import torch
import torch.nn as nn

# 単純なTensorFlowモデルを定義
class SimpleModel(tf.keras.Model):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

# 1. データ読み込み
# 簡単なTensorFlowモデルを作成・学習
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784) / 255.0
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
tf_model = SimpleModel()
tf_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
tf_model.fit(x_train, y_train, epochs=5)

# 2. 形式変換
# 対応するPyTorchモデルを定義
class PyTorchModel(nn.Module):
    def __init__(self):
        super(PyTorchModel, self).__init__()
        self.dense1 = nn.Linear(784, 64)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.dense1(x)
        x = self.relu(x)
        x = self.dense2(x)
        return x

# TensorFlowモデルの重みをPyTorchモデルに転送
pytorch_model = PyTorchModel()
with torch.no_grad():
    pytorch_model.dense1.weight = nn.Parameter(torch.tensor(tf_model.layers[0].get_weights()[0].T))
    pytorch_model.dense1.bias = nn.Parameter(torch.tensor(tf_model.layers[0].get_weights()[1]))
    pytorch_model.dense2.weight = nn.Parameter(torch.tensor(tf_model.layers[1].get_weights()[0].T))
    pytorch_model.dense2.bias = nn.Parameter(torch.tensor(tf_model.layers[1].get_weights()[1]))

# 3. モデル保存
# 変換されたPyTorchモデルを保存
torch.save(pytorch_model.state_dict(), 'pytorch_model.pth')

上記コードでは、MNISTデータセット上で学習されたTensorFlowモデルSimpleModelを定義し、対応するPyTorchモデルPyTorchModelを構築します。TensorFlowモデルのパラメータをPyTorchモデルに移行し、最終的に保存することで、PyTorchでの推論が可能になります。

モデル変換ツール

以下は、さまざまなフレームワーク間のモデル移行を可能にする変換ツールの一覧です：

規格化変換

ここでは、ONNXを例として規格化変換の仕組みを紹介します。

ONNXの概要

ONNX（Open Neural Network Exchange）は、機械学習モデルを格納するためのオープンなファイルフォーマットです。これにより、PyTorchやMXNetなどの異なるAIフレームワーク間でモデルを共有し、相互運用することが可能になります。

ONNXはマイクロソフト、アマゾン、Meta、IBMなどの企業によって開発され、GitHub上でオープンソースとして公開されています。現在、ONNXモデルをロードし推論を行うことをサポートしているフレームワークには、Caffe2、PyTorch、MXNet、ML.NET、TensorRT、Microsoft CNTKなどがあり、TensorFlowも非公式にONNXをサポートしています。

各AIフレームワークは独自のグラフ表現とAPIを持つため、フレームワーク間でのモデル変換は複雑です。また、トレーニング速度、複雑なネットワーク構造のサポート、モバイル端末での推論など、それぞれのフレームワークは異なる最適化を施しています。ONNXは共通の計算グラフ表現を提供し、開発者がどの段階でも最適なフレームワークを選択できるようにします。

ONNXは拡張可能な計算グラフモデル、ビルトイン演算子(OP)、標準データ型を定義します。計算グラフはノードのリストとして定義され、有向非巡回グラフ(DAG)として構築されます。各ノードは入力と出力を持ち、OPと呼ばれます。これは汎用的な計算グラフであり、さまざまなAIフレームワークで作成されたグラフを変換できます。

PyTorchから ONNXへの変換例

前節で保存したPyTorchモデルpytorch_model.pthを読み込み、torch.onnx.export()を使用してONNX形式に変換します：

x = torch.randn(1, 784)

# ONNX形式へエクスポート
with torch.no_grad():
    torch.onnx.export(
        pytorch_model,
        x,
        "pytorch_model.onnx",
        opset_version=11,
        input_names=['input'],
        output_names=['output']
    )

上記コードが正常に実行されれば、pytorch_model.onnxというファイルが生成されます。以下のように確認できます：

import onnx 

onnx_model = onnx.load("pytorch_model.onnx") 
try: 
    onnx.checker.check_model(onnx_model) 
except Exception: 
    print("Model incorrect") 
else: 
    print("Model correct")

onnx.loadはONNXモデルを読み込み、onnx.checker.check_modelはモデル形式の正当性を検証します。問題がない場合は「Model correct」と表示されます。

Netron（オープンソースのモデルビジュアルツール）を使用してONNXモデルを視覚化できます。入力または出力クリックで、モデルバージョンや入出力名・データ型を確認できます。演算子ノードをクリックすると、属性・構造・重み情報を確認できます。

モデル変換の一般的なプロセス

1. 計算グラフ生成：AIフレームワークは静的グラフ形式で計算グラフを生成します。一般的にはASTベースの変換とTraceベースの方法があります。
   • ASTベースの変換：動的グラフを静的グラフに自動的に変換します。ソースコードを解析し、抽象構文木を変換することで、制御フローの欠落を防ぎ、静的グラフモデルが元の動的グラフと同等の動作をするようにします。
   • Traceベース：動的グラフモードで実行し、実行された演算子を記録し、その順序に基づいて静的グラフを構築します。これにより、動的グラフの動作を正確に再現できます。
2. 共通演算子の対応：モデル内の共通演算子がターゲットフレームワークで実装されていることを確認します。カスタム演算子については、専用の変換ロジックを実装するか、等価な共通演算子の組み合わせに置き換えます。
3. 中間形式への変換：モデルを推論エンジンのカスタムIR形式に変換します。IRには計算グラフ、演算子、パラメータが含まれ、モデル変換の柔軟性と効率性を高めます。
4. モデルの出力：推論エンジンのIR形式からモデルファイルを出力・保存し、後続の推論処理に使用します。

変換プロセスでは、ソースとターゲットフレームワークの演算子互換性、テンソル形式の違い、および変換後のモデルの精度保持と最適化が重要です。