AWS 深層学習 AMI に ONNX が含まれ、深層学習フレームワーク間でのモデルの移植性が向上

Ubuntu および Amazon Linux 用の AWS 深層学習 AMI (DLAMI) に完全に設定済みの Open Neural Network Exchange (ONNX) がプリインストールされることになり、深層学習フレームワーク間でのモデルの移植性が向上しました。このブログ記事では、ONNX を紹介し、DLAMI で ONNX を使用してフレームワーク間でモデルを移植する方法を示します。

ONNX とは

ONNX は、オープンソースライブラリであり、シリアライゼーションフォーマットを使って深層学習モデルをエンコードおよびデコードします。ONNX は、ニューラルネットワークの計算グラフのフォーマットと、ニューラルネットワークアーキテクチャで使用される演算子の広範なリストを定義します。ONNX は、Apache MXNet、PyTorch、Chainer、Cognitive Toolkit、TensorRT などの一般的な深層学習フレームワークですでにサポートされています。普及しているツールで ONNX のサポートが拡大することにより、機械学習の開発者は、ツールの違いを超えてモデルを移動し、必要な作業に最適なツールを選択することができるようになります。

Chainer モデルを ONNX にエクスポートする

それでは、Chainer モデルを ONNX ファイルにエクスポートする手順を見てみましょう。

まず、Ubuntu または Amazon Linux で DLAMI のインスタンスを起動します。以前に起動したことがない場合は、DLAMI を使い始める方法を説明しているこの素晴らしいチュートリアルをご覧ください。

SSH 経由で DLAMI に接続したら、DLAMI に設定済みでプリインストールされている Chainer Python 3.6 Conda 環境を有効にします。また、この環境には ONNX と、Chainer に対する ONNX サポートを追加するアドオンパッケージである onnx-chainer もプリインストールされ、設定済みであることに注意してください。

$ source activate chainer_p36

次に、Python シェルを起動し、以下のコマンドを実行して、物体認識のための VGG-16 畳み込みニューラルネットワークを読み込み、ニューラルネットワークを ONNX ファイルにエクスポートします。

import numpy as np
import onnx_chainer
from chainercv.links import VGG16

# This will download a pre-trained model, and load it as a Chainer model
model = VGG16(pretrained_model='imagenet')

# Creating synthetic input and using it to export the model to ONNX
x = np.zeros((1, 3, 224, 224), dtype=np.float32)
out = onnx_chainer.export(model, x, filename='vgg16.onnx')

こうして、わずか数行のコードで、Chainer モデルを ONNX 形式にエクスポートし、現在のディレクトリにファイルを保存することができました。

ONNX モデルを MXNet にインポートする

Chainer モデルを ONNX にエクスポートしたので、次は、このモデルを MXNet にインポートして推論を実行する方法を見てみましょう。

まず、ONNX および MXNet 1.2.1 があらかじめインストールされている、DLAMI の MXNet Python 3.6 Conda 環境を有効にします。MXNet 1.2 は、ONNX モデルを MXNet にインポートするために使用する ONNX インポート API を導入します。

$ source deactivate
$ source activate mxnet_p36

次に、Python シェルを起動し、以下のコマンドを実行して、さきほど Chainer からエクスポートした ONNX モデルを読み込みます。

from mxnet.contrib import onnx as onnx_mxnet
sym, arg_params, aux_params = onnx_mxnet.import_model("vgg16.onnx")

これで完了です。ONNX モデルを MXNet に読み込み、シンボルグラフとパラメータを使用できるようにしました。新たに読み込んだモデルで推論を続けましょう。

最初に画像をダウンロードし、モデルをトレーニングした ImageNet クラスのラベルをダウンロードして、物体認識モデルをテストします。

import mxnet as mx
mx.test_utils.download('https://s3.amazonaws.com/onnx-mxnet/dlami-blogpost/hare.jpg')
mx.test_utils.download('http://data.mxnet.io/models/imagenet/synset.txt')
with open('synset.txt', 'r') as f:
    labels = [l.rstrip() for l in f]

入力画像は次のようになります。

次に、画像を読み込み、モデルに必要な入力テンソルの形状に合ったテンソルに前処理します。

import matplotlib.pyplot as plt
import numpy as np
from mxnet import nd

image = plt.imread("hare.jpg")
image = np.expand_dims(np.transpose(image, (2,0,1)),axis=0).astype(np.float32)
input = nd.array(image)

これで、MXNet モジュールを初期化してバインドする準備が整いました。

# Input name is the model’s input node name, and is defined by the exporting library
input_name = sym.list_inputs()[0] 
data_shapes = [(input_name, input.shape)]
# Initialize and bind the Module
mod = mx.mod.Module(symbol=sym, context=mx.cpu(), data_names=[input_name], label_names=None)
mod.bind(for_training=False, data_shapes=data_shapes, label_shapes=None)
mod.set_params(arg_params=arg_params, aux_params=aux_params)

そして最後に、推論を実行して最高の確率とクラスを出力します。

mod.forward(mx.io.DataBatch([input]))

probabilities = mod.get_outputs()[0].asnumpy()[0]
max_probability = np.max(probabilities)
max_class = labels[np.argmax(probabilities)]

print('Highest probability=%f, class=%s' %(max_probability, max_class))

出力が、モデルの予測を示します。このモデルは、画像は 97.9% の確率で野ウサギであると予測しています！

最高確率=0.979566、クラス=n02326432 野ウサギ

結論および深層学習 AMI の開始方法

このブログ記事では、DLAMI で ONNX を使用してフレームワーク間でモデルを移植する方法を学びました。ONNX によって移植が可能になることで、新しいモデルをトレーニングしたり、事前にトレーニングしたモデルを微調整したり、推論やモデル提供を実行したりするなど、扱っているタスクに最適なツールを選択することができるようになります。

入門チュートリアルを使用すると、AWS 深層学習 AMI を速やかに始めることができます。その他のチュートリアル、リソース、リリースノートについては、DLAMI ONNX チュートリアル、ならびに開発者ガイドをご覧ください。最新の AMI が AWS Marketplace で利用可能になりました。また、フォーラムをサブスクライブして、新しい開始の通知を受け取ったり、質問を投稿したりすることもできます。

ブログ投稿者について

Anirudh Acharya は AWS 深層学習のソフトウェア開発エンジニアです。彼は、誰でも AI を使えるようにするための深層学習システムとツールキットの構築に取り組んでいます。読書、自転車を楽しむことが趣味です。

Hagay Lupesko は AWS 深層学習のエンジニアリングリーダーです。開発者やサイエンティストがインテリジェントアプリケーションを構築できるようにする深層学習システムを担当しています。そのかたわらで、読書、ハイキング、家族との時間を楽しんでいます。

Amazon Web Services ブログ