第二回ONNX.jsの備忘録(Pythonで作ったモデルを扱う)

はじめに

ONNX.jsの練習用にいろいろやっていくための備忘録
今回は簡単なモデルを作り，それをONNX形式にエクスポートしてみる．

なお前回の記事はこちらONNXjsでResNetを動かす

今回の目的

ONNX.jsの使い方はなんとなく，ぼんやり，それとなくわかったので，今回はPythonでモデルを作って，そいつをONNX形式に出力する．

例によって先人のすごい方々が記事を残してくれている．
今回も大いに参考にさせていただきました．
初めてのONNX
ONNX形式のモデルを扱う

公式のドキュメントもかなりわかりやすい．
ONNX PythonAPI

学習済みのモデル一覧というものもあるようだ．

Pythonでモデルを作る

いろいろと下準備

まずはONNXのPython用のモジュールがないと始まらないのでインストール．pipでインストールできる．

pip install onnx

公式にも詳しくインストール方法が記載されている．
NumpyやらPytorchなども必要になってくるが，ここでは割愛．足りないものがあれば都度追加していく．

環境構築していない場合Anacondaでやってしまうのがラク．

モデル作成

今回はとにかく簡単化するため，2入力1出力のモデルを作成する．
ちなみに1入力で試したところonnxにダメだと怒られた(Node has input size 1 not in range)．

import onnx
import onnx.helper as helper
from onnx import AttributeProto, TensorProto, GraphProto, checker


# 2入力，1中間，1出力のモデルを作る，結果をバイナリ形式とテキスト形式で保存
# 参考：ONNX公式                https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
# 参考：初めてのONNX　           https://qiita.com/natsutan/items/6670bd518005152f654a
# 参考：ONNX形式のモデルを扱う　   https://qiita.com/nuka137/items/905dddfc62ba32357aa1


# 入力のテンソルを作成
# どうやら入力が1だとエラーが出るっぽい（Node  has input size 1 not in range）
# make_tensor_value_infoはValueInfoProtoを作成する関数
# make_tensor_value_infoの引数は(name[ValueInfoProtoの名前], elem_type[データの型], shape[テンソルの形], doc_string[テキスト], shape_denotation)
Input_X1 = helper.make_tensor_value_info('Input_X1', TensorProto.FLOAT, [3, 2])
Input_X2 = helper.make_tensor_value_info('Input_X2', TensorProto.FLOAT, [3, 2])


# 出力のテンソルを作成
Output_Y1 = helper.make_tensor_value_info('Output_Y1', TensorProto.FLOAT, [3, 2])

# 計算結果を引数にしてノードを追加していく
node_def = helper.make_node(
    'Pad',                          # name(op_type operatorの名前)
    ['Input_X1', 'Input_X2'],       # inputs(入力の文字列をリストで指定)
    ['Output_Y1'],                  # outputs(出力の文字列をリストで指定)
    mode = 'constant',              # attributes(属性の指定)
)

graph_def = helper.make_graph(
    [node_def],                     # nodes(NodeProtoをリストで指定)
    'test-model',                   # name(グラフの名前)
    [Input_X1, Input_X2],           # inputs(入力に使うValueInfoProtoをリストで指定)
    [Output_Y1],                    # outputs(出力に使うValueInfoProtoをリストで指定)
)

# モデルを構築
model_def = helper.make_model(
    graph_def, 
    producer_name='Example',
    )

# モデルの出力ファイル名を設定
out_path = "model_Example.onnx"

# 作成したモデルをファイルへ保存（バイナリ形式）
with open(out_path, "wb") as f:
    f.write(model_def.SerializeToString())

# 作成したモデルをテキスト形式でも保存
with open(out_path + ".txt", "w") as f:
    print(model_def, file=f)


print('The model is checked!')

というわけでモデルを作って，ファイルに出力までを完了．
出力されたテキストファイルの中身はこんなかんじ．

ir_version: 7
producer_name: "Example"
graph {
  node {
    input: "Input_X1"
    input: "Input_X2"
    output: "Output_Y1"
    op_type: "Pad"
    attribute {
      name: "mode"
      s: "constant"
      type: STRING
    }
  }
  name: "test-model"
  input {
    name: "Input_X1"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  input {
    name: "Input_X2"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  output {
    name: "Output_Y1"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
}
opset_import {
  version: 14
}

とりあえずは正常にできているようだ．
このままではわかりにくいのでNetronでモデルを可視化してみる．

なんともお粗末なモデルであるが，とりあえず想定していたものは出来上がり．
ひとまずこれにて，Pythonでモデルを作ってONNXに出力するという当初の目的は達成である．
次回はもう少しパラメータの設定など細かいところを勉強しつつ，込み入ったモデルを作ってみたい（願望）．