Apache MXNet で ONNX をサポート

AWS は ONNX-MXNet の利用開始を発表しました。これは Open Neural Network Exchange (ONNX) ディープラーニングモデルを Apache MXNet にインポートするためのオープンソース Python パッケージです。MXNet は充実した機能を備えたスケーラブルなディープラーニングフレームワークで、Python、Scala、R といった人気の言語に対し API を提供します。MXNet で ONNX 形式をサポートすることで、開発者は PyTorch、Microsoft Cognitive Toolkit、Caffe2 など、他のフレームワークを使用してモデルを構築したりトレーニングすることができます。また、高度に最適化されたスケーラブルなエンジンの MXNet を使用した推論に対し、こうしたモデルを MXNet にインポートすることもできます。

AWS が ONNX 形式に貢献できることを大変喜ばしく思っています。Facebook、Microsoft、そしてディープラーニングコミュニティと協力し、ディープラーニングのユーザーが利用しやすい便利なものにすべく、ONNX の開発に取り組みます。

ONNX とは

ONNX はディープラーニングモデルをエンコードするためのオープンソース形式です。ONNX はニューラルネットワークの計算グラフ、グラフ内で使用される演算子の広範なリストの形式を定義します。拡大中のフレームワークリスト、ハードウェアベンダー、ディープラーニングの開発を手掛ける開発者などにサポートされている ONNX は、容易にフレームワーク間を移動し、目の前の課題に最適なフレームワークを選別することができます。

クイックスタート

今回は ONNX-MXNet を使用して MXNet に ONNX-MXNet をインポートする方法、そして推論用にインポートしたモデルを使用する方法をご紹介します。これにより、MXNet の最適化した実行エンジンのメリットを活用することができます。

ステップ 1: インストール

まず、ONNX repo の手順に従い、ONNX をインストールします。

次に ONNX-MXNet パッケージをインストールします。

$ pip install onnx-mxnet

ステップ 2: ONNX モデルを準備してインポート

この例では、イメージの空間分解能を高めるように設計されている Super Resolution モデルのインポートを行います。このモデルは PyTorch で構築とトレーニングが行われ、PyTorch の ONNX エクスポート API を使用して ONNX にエクスポートされています。モデル設計の詳細については PyTorch の例をご覧ください。

Super Resolution ONNX モデルを作業ディレクトリにダウンロードします。

$ wget https://s3.amazonaws.com/onnx-mxnet/examples/super_resolution.onnx

ステップ 3: ONNX モデルを MXNet にインポート

ONNX モデルのファイルの準備ができたので、ONNX-MXNet インポート API を使用して MXNet にインポートしてみましょう。Python シェルで次のコードを実行します。

import onnx_mxnet
sym, params = onnx_mxnet.import_model('super_resolution.onnx')

Python ランタイムで 2 つのインスタンスを作成しました。 sym – モデルのシンボリックグラフと params – モデルのウェイトです。これで ONNX モデルのインポートが完了し、標準の MXNet モデルができました。

ステップ 4: 推論用の入力を準備

推論用の入力イメージを準備します。次の手順で例のイメージをダウンロードし、モデルの予期された入力形式にサイズを変更してから、numpy 配列に変換します。

シェルコンソールから、入力イメージの例を作業ディレクトリにダウンロードします。

$ wget https://s3.amazonaws.com/onnx-mxnet/examples/super_res_input.jpg

次に入力イメージをロードし事前処理できるように Pillow、Python Imaging Library をインストールします。

$ pip install Pillow

Python シェルからコードを実行して MXNet NDArray 形式のイメージを準備します。

import numpy as np
import mxnet as mx
from PIL import Image
img = Image.open("super_res_input.jpg").resize((224, 224))
img_ycbcr = img.convert("YCbCr")
img_y, img_cb, img_cr = img_ycbcr.split()
test_image = mx.nd.array(np.array(img_y)[np.newaxis, np.newaxis, :, :])

ステップ 5: MXNet モジュールを作成

MXNet モジュール API を使用してモジュールを作成し、バインドしてロード済みウェイトを指定します。
ONNX-MXNet インポート API は、入力レイヤー「input_0」を指定することに注意してください。これは初期化中およびモジュールのバインドに使用します。

mod = mx.mod.Module(symbol=sym, data_names=['input_0'], label_names=None)
mod.bind(for_training=False, data_shapes=[('input_0',test_image.shape)])
mod.set_params(arg_params=params, aux_params=None)

ステップ 6: 推論の実行

MXNet モジュールのロード、バインド、トレーニング済みウェイトが完了したので、推論実行の準備が完了しました。単一の入力バッチの準備を行い、ネットワークを介してフィードフォワードします。

from collections import namedtuple
Batch = namedtuple('Batch', ['data'])
mod.forward(Batch([test_image]))
output = mod.get_outputs()[0][0][0]

ステップ 7: 結果を検証

では、Super Resolution イメージで実行した推論の結果を検証してみましょう。

img_out_y = Image.fromarray(np.uint8((output.asnumpy().clip(0, 255)), mode='L'))
result_img = Image.merge(
"YCbCr", [
        	img_out_y,
        	img_cb.resize(img_out_y.size, Image.BICUBIC),
        	img_cr.resize(img_out_y.size, Image.BICUBIC)
]).convert("RGB")
result_img.save("super_res_output.jpg")

入力イメージと結果の出力イメージです。ご覧のように、モデルはイメージの空間分解能を 256 X 256 から 672 X 672 に上げることができました。

入力イメージ	出力イメージ

次のステップ?

AWS は ONNX パートナーやコミュニティと協力して ONNX のさらなる開発、より便利な演算子の追加、エクスポートを含む ONNX-MXNet の拡張と演算子の範囲を拡大するように努めています。Apache MXNet コミュニティと協力し、今後 ONNX を MXNet コア API に取り入れられるように取り組んでいきます。