ONNXを試す

前回はONNXとは何なのか、インストール方法を書いたので、 今回は実際に使ってみた、という記事です。

ONNXデータのインポート

まずはデータを取得するところからです。 サンプルとしてONNXのモデルのレポジトリに記載されているリンクから好きなモデルのファイルを取得します。ここではinception v1モデルのonnxファイルをダウンロードしたいと思います。一応以下のコマンドからも取得できます。

$ wget https://s3.amazonaws.com/download.onnx/models/opset_9/inception_v1.tar.gz

これを解凍すると、

inception_v1/
  |-model.onnx
  |-test_data_0.npz
  |-test_data_1.npz
  |-test_data_2.npz
  |-test_data_set_0/
  |-test_data_set_1/
  |-test_data_set_2/

このようなディレクトリ構成になります。 とりあえずモデルの情報を見てみましょう。

import onnx
model_path = "./model.onnx"
model = onnx.load(model_path)
for node in model.graph.node:
   print(node)

これだけです。実行結果は

input: "inception_4c/output_1"
input: "inception_4d/1x1_w_0"
input: "inception_4d/1x1_b_0"
output: "inception_4d/1x1_1"
name: ""
op_type: "Conv"
attribute {
  name: "strides"
  ints: 1
  ints: 1
  type: INTS
}
attribute {
  name: "pads"
  ints: 0
  ints: 0
  ints: 0
  ints: 0
  type: INTS
}
attribute {
  name: "kernel_shape"
  ints: 1
  ints: 1
  type: INTS
}

こういう感じの構造がいっぱい出てきます。入出力の指定、演算の指定、ストライドやパディング、カーネルのサイズなどの情報を得られるのでこのonnxファイルをもとにどのような構造のモデルなのかを解析することができますね。

Netronを試す

ONNXのモデルを上記のようにインポートできたら次はNetronというVisualizerを起動させてみます。webとGUIアプリのどちらからも選べます。今回はwebアプリを使ってみます。 https://lutzroeder.github.io/netron/ にアクセスして、"Open Model"でダウンロードしてきたonnxファイルを選択します。 そうするとしばらく待つと次のような図が出てきて、きちんとモデルに関する情報が抽出できていることがわかります。 それぞれの層のブロックをクリックすると計算に用いるカーネルサイズやパディング、重みの値なども見ることができます。

ONNXへモデルをエクスポート

onnxファイルからモデルを取得することはできたので、次は自分で作成したモデルをエクスポートしてみます。

import onnx                                                                                                         
import onnx.helper as helper

export_onnx_path = "sample_model.onnx"

# define input
input_tensor = []
input_tensor.append(helper.make_tensor_value_info("input_1", onnx.TensorProto.FLOAT, [24, 24, 3]))
input_tensor.append(helper.make_tensor_value_info("input_2", onnx.TensorProto.FLOAT, [24, 24, 3]))

# define output
output_tensor = []
output_tensor.append(helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [24, 24, 3]))

# define computational graph
nodes = []
nodes.append(helper.make_node("Add", ["input_1", "input_2"], ["output"]))

graph = helper.make_graph(nodes, "Sample graph", input_tensor, output_tensor)

model = helper.make_model(graph)

with open(export_onnx_path, "wb") as f:
    f.write(model.SerializeToString())

割とシンプルな書き方で何がしたいかがわかりやすいでしょうか？ どのような計算グラフを作成したかを想像してもらった上で、 このスクリプトで作成したonnxファイルを改めて先程のNetronで見てみると、 こんな感じです。 こんな小さなグラフでも可視化すると達成感があります。 実際には自分で計算グラフを構成するようなことは少なくて、ある機械学習フレームワークで作成した学習モデルを他のフレームワークに移植するというために使うことが多いです。 そのためのチュートリアルのレポジトリもあるので見てみたいと思います。