ONNX-Chainer紹介

## ONNXとは

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

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

基本的にONNX形式での出力を行うフレームワークは[onnx/onnx](https://github.com/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](https://nvidianews.nvidia.com/news/nvidia-gpu-cloud-now-available-to-hundreds-of-thousands-of-ai-researchers-using-nvidia-desktop-gpus)

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

## ONNX-Chainerのインストール

ONNX-Chainerは、[Chainer](https://github.com/chainer/chainer/)と[ONNX](https://github.com/onnx/onnx)という2つのパッケージに依存しています。ONNX-Chainerはpipからインストールすることができます。

```bash
pip install onnx-chainer
```

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

## 使い方

ONNX-Chainerは現在、50種類の関数の出力に対応しています（PyTorchの最新版0.3.0では、36種類の関数に対応）。こちらの[README](https://github.com/chainer/onnx-chainer#supported-functions)にその一覧があります。ここにある関数のみで構成されたネットワークであれば、その構造とパラメータを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形式で保存するコードを示します。

```python
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](https://github.com/onnx/onnx-caffe2)をインストールしておきます。これらは外部ライブラリなので導入方法などは公式情報を参照してください。

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

```python
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形式のモデルに変換されたものを保存できます。

```python
# 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](https://github.com/onnx/onnx/blob/master/docs/Operators.md)。あとは、出力したいChainerのFunctionNodeの仕様をコードを見て理解してパラメータ名などを対応させるだけです。

まずここから対象のChainerクラスをさがします：[chainer/functions](https://github.com/chainer/chainer/tree/v3/chainer/functions)。次に、[Operator Schemas](https://github.com/onnx/onnx/blob/master/docs/Operators.md)にある対応する関数の定義を探します。

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

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

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

| 引数名 | 説明 |
|:------|:----|
| `func` | 変換したいChainerの`FunctionNode`オブジェクトが渡される |
| `input_names` | その`FunctionNode`に入力された`Variable`の[識別値](https://docs.python.jp/3/library/functions.html#id)が並んだもの |
| `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](https://github.com/chainer/onnx-chainer/blob/master/onnx_chainer/functions/connection/convolution_2d.py#L6-L36)の定義などを見るとわかりやすいかと思います。

もし新しい関数のサポートを追加してくださった場合は、[chainer/onnx-chainer](https://github.com/chainer/onnx-chainer/pulls)までPRを送ってくださると嬉しいです。

## おわりに

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