ONNX-Chainer紹介

ONNXとは

ONNXは、Open Neural Network Exchangeの略で、Deep Learningモデルを表現するためのフォーマットです。Chainer, MXNet, Caffe2などいろいろなフレームワークがありますが、各フレームワークがこのONNXというフォーマットでのモデルの保存・読み込みに対応することで、同じモデルを異なるフレームワーク間で受け渡し合うことができるようになります。

ONNXでは、Protocol Bufferというシリアライズフォーマットを使って、ニューラルネットワークでよく使われる様々なオペレータが定義されています。ONNX形式での出力が可能なオペレータの種類は、こちらに一覧されています：Operator Schemas。これとは別に、各フレームワークがどのようなオペレータの入出力に対応しているかがモデルを受け渡す際には重要になります。例えば、Chainerが畳み込みレイヤの出力に対応していても、受け渡したい先のフレームワークが畳み込みレイヤの読み込みに対応していなければ、畳み込みを含むネットワークの受け渡しはできません。

基本的にONNX形式での出力を行うフレームワークはonnx/onnxというライブラリを使って、フレームワーク内のネットワーク構造とパラメータをONNX形式で出力する機能を追加する必要があります。また、ONNX形式を読み込む側のフレームワークも同様に、このライブラリを使ってONNX形式のファイルの読み込みを行う機能を追加する必要があります。

現在主要なフレームワークでのONNX対応状況は2017年12月17日現在、以下のようになっています。

Framework	ONNX Import	ONNX Export
Chainer	x	o
PyTorch	x	o
MXNet	o	x
Caffe2	o	o
CNTK	o	o
TensorFlow	x	x

また、TensorRTもONNX対応を表明しています：NGC Expands Further, with NVIDIA TensorRT Inference Accelerator, ONNX Compatibility, Immediate Support for MXNet 1.0

本記事では、chainer/onnx-chainerを使ってこのONNX形式のファイルにChainerで記述したモデルを出力する方法と、新しくonnx-chainerにサポートするオペレータを追加する手順を説明します。

ONNX-Chainerのインストール

ONNX-Chainerは、ChainerとONNXという2つのパッケージに依存しています。ONNX-Chainerはpipからインストールすることができます。

pip install onnx-chainer

上記コマンドで現時点ではONNX v1形式に対応した最新版のONNX-Chainerがインストールされます。

使い方

ONNX-Chainerは現在、50種類の関数の出力に対応しています（PyTorchの最新版0.3.0では、36種類の関数に対応）。こちらのREADMEにその一覧があります。ここにある関数のみで構成されたネットワークであれば、その構造とパラメータをONNX形式で出力することができます。方法はとても簡単です。

ChainerはDefine-by-Runというスタイルを採用しているため、ネットワーク構造は実際にデータを渡してフォワード計算を行うまで決定されていません。そのため、ONNX-Chainerではモデルを表すChainオブジェクトと、そこに流し込むための想定される形式・型のデータを渡してやる必要があります。ONNX-Chainerは、渡されたChainオブジェクトとデータを使って一度フォワード計算を実行しネットワーク構造を決定します。その際、渡してやるChainオブジェクトには__call__メソッドが用意されていることが前提となります。そして、これらの処理を行いネットワーク構造を決定してONNX形式への変換を行うonnx_chainer.exportの第二引数に渡されるものは、そのまま第1引数で渡されたChainオブジェクトの__call__に渡されます。実例を見たほうがわかりやすいと思うので、以下にChainerが標準で用意しているVGG16のPre-trainedモデルをONNX形式で保存するコードを示します。

import numpy as np
import chainer
import chainer.links as L
import onnx_chainer

model = L.VGG16Layers()

# ネットワークに流し込む擬似的なデータを用意する
x = np.zeros((1, 3, 224, 224), dtype=np.float32)

# 推論モードにする
chainer.config.train = False

onnx_model = onnx_chainer.export(model, x, filename='VGG16.onnx')

modelはChainオブジェクトです。xは想定される入力データを実際にnumpy.ndarrayオブジェクトとして用意したものです。onnx_chainer.export()メソッドにこれらが渡されると、内部でmodel(x)が実行され、その出力をもとにネットワーク構造が取り出されます。その他の引数や使用等は、以下のようになっています。

def export(model, args, filename=None, export_params=True,
           graph_name='Graph', save_text=False):

    Export function for chainer.Chain in ONNX format.
    This function performs a forward computation of the given Chain, model, by passing the given arguments args directly. It means, the output Variable object `y` to make the computational graph will be created by:

    y = model(*args)

    Args:
        model (~chainer.Chain): The model object you want to export in ONNX format. It should have __call__ method because the second argument args is directly given to the model by the () accessor.
        args (list or dict): The arguments which are given to the model directly.
        filename (str or file-like object): The filename used for saving the resulting ONNX model. If None, nothing is saved to the disk.
        export_params (bool): If True, this function exports all the parameters included in the given model at the same time. If False, the exported ONNX model doesn't include any parameter values.
        graph_name (str): A string to be used for the "name" field of the graph in the exported ONNX model.
        save_text (bool): If True, the text format of the output ONNX model is also saved with ".txt" extension.

    Returns:
        A ONNX model object.

ONNX出力したChainerモデルをCaffe2で実行

事前にCaffe2本体とonnx/onnx-caffe2をインストールしておきます。これらは外部ライブラリなので導入方法などは公式情報を参照してください。

先程と同様にまずはONNX-Chainerを使ってChainerのChainオブジェクトをONNX形式に変換します。onnx_chainer.export()は、ONNXモデルのオブジェクトを返すので、これを受け取っておき、onnx_caffe2のrun_modelに渡せばそれをCaffe2を使って走らせることができます。

import chainer
import chainer.links as L
import numpy as np

from onnx_caffe2.backend import Caffe2Backend
from onnx_caffe2.backend import run_model
from onnx_caffe2.helper import save_caffe2_net

import onnx_chainer

# Instantiate a Chainer model (Chain object)
model = L.VGG16Layers()

# Prepare a dummy input
x = np.random.randn(1, 3, 224, 224).astype(np.float32)

# Do not forget setting train flag off!
chainer.config.train = False

# Export to ONNX model
onnx_model = onnx_chainer.export(model, x)

# Run the model with Caffe2
caffe2_out = run_model(onnx_model, [x])[0]

ONNX出力したChainerのモデルをCaffe2形式で保存

onnx_caffe2の機能を使うと、明示的にONNXモデルをCaffe2のモデルに変換したのち、それをファイルに保存することができます。上のコードでonnx_modelを作成した後、以下のようなコードを実行すれば、Caffe2形式のモデルに変換されたものを保存できます。

# Convert ONNX model to Caffe2 model
init_net, predict_net = Caffe2Backend.onnx_graph_to_caffe2_net(
    onnx_model.graph, device='CPU')

# Save the Caffe2 model to disk
init_file = "./vgg16_init.pb"
predict_file = "./vgg16_predict.pb"
save_caffe2_net(init_net, init_file, output_txt=False)
save_caffe2_net(predict_net, predict_file, output_txt=True)

Caffe2はiOSやAndroidなどのモバイルデバイス上での実行もサポートしているので、これで理論上はChainerで試行錯誤して学習したモデルを、モバイルデバイスにデプロイできるようになった、と言えます。

対応レイヤを追加する

ONNX exportに対応したChainerの関数を増やす場合は、まずONNX自体がサポートしていることを確認してください：Operator Schemas。あとは、出力したいChainerのFunctionNodeの仕様をコードを見て理解してパラメータ名などを対応させるだけです。

まずここから対象のChainerクラスをさがします：chainer/functions。次に、Operator Schemasにある対応する関数の定義を探します。

次に、onnx-chainer/mapping.pyに関数名の対応を記述します。operatorsというディクショナリがあるので、ここにChainer側の関数名をキー、ONNX側のオペレータ名を値にした要素を追加します。

次に、onnx-chainer/functions以下にChainerのソースツリーと対応した名前のディレクトリが並んでいるので、自分が追加したい関数が属するディレクトリの下に、Chainerでのファイル名と対応した名前のファイルを追加します。そこにconvert_[変換したいChainerのFunctionNodeのクラス名]というメソッドを一つ書いておき、そのメソッド名をonnx-chainer/functions/init.pyの適切な行に追加して、onnx_chainer.export()がChainerの関数名を元に自動的に変換する関数を見つけられるようにしておきます。

最後に、実際に変換を行う部分を先程のconvert_[変換したいChainerのFunctionNodeのクラス名]（例えば、convert_Convolution2DFunction）に書いていきます。この関数は、以下の5つを引数にとるように書いてください。

引数名	説明
func	変換したいChainerのFunctionNodeオブジェクトが渡される
input_names	そのFunctionNodeに入力されたVariableの識別値が並んだもの
param_names	ChainerでChain.named_params()を呼び出した際に返されるパラメータの名前のリスト（/l1/W, /l1/bなど）
parameters	Chain.params()が返すParameterオブジェクトから抜き出したarray要素（numpy.ndarray）
input_tensors	ネットワーク自体への入力をONNXのTensorValueInfoというクラスのオブジェクトに変換したもの

ONNXのフォーマットへ変換を行う際に、これらの全ては使わないかもしれませんが、変換を行う関数名は変換元のChainerのFunctionNodeのクラス名から自動的に決定されたのち、どの関数に対する変換を行う関数かによらず、同じデータが引き渡されるように現状ではなっています。

convert_[変換したいChainerのFunctionNodeのクラス名]の中では、先程対応を追加したmapping.pyの中のoperatorsディクショナリを読み込んで、対応するONNX側のオペレータ名を取得し、onnxライブラリのヘルパー関数make_nodeにそのオペレータ名、そのレイヤへの入力の名前のリスト、出力のリスト、そしてレイヤごとに異なるattributeをキーワード引数として渡す形で変換が行なえます。

詳しくはconvert_Convolution2DFunctionの定義などを見るとわかりやすいかと思います。

もし新しい関数のサポートを追加してくださった場合は、chainer/onnx-chainerまでPRを送ってくださると嬉しいです。

おわりに

Chainer Advent Calendar 14日目の記事で @okdshin さんがC++でCPUのみでonnx-chainerで出力したONNX形式のDNNモデルの推論（Inference）をするライブラリ「Instant」を作ったという記事を書いてくださいました。このInstantを使うと、Chainerで試行錯誤し学習したモデルをC++のプロダクトから読み込んで実行するまでが容易になります。