ONNXのValueInfoProtoにinitializerで値を設定する

これの続き。自力でMLPのONNXファイルを作りたい。

Reshapeのパラメータ（reshape後のshape)が、値が設定されていないのでグラフの入力につながないとエラーになる。

こちらを見ながら、initializerで値を設定する。
https://github.com/onnx/onnx/blob/master/docs/IR.md

initializerはValueInfoProtoではなく、GraphProtoに紐づくのでmake_graphの引数で指定する。値はnumpy.helperを使い、numpy arrayからTensorProtoを作る。TensorProtoにValueInfoProtoと同じ名前を付けると、make_graph時に紐づけてくれる。

import numpy
from onnx import helper, numpy_helper

# numpy array から TensorProtoを作る
input_shape_init = numpy_helper.from_array(numpy.array([784], dtype=numpy.int64))
input_shape_init.name = 'input_shape'

これで、Graphの入力からReshapeのパラメータが消えました。

全ソースはこう。次はレイヤーを書さえてMLPにする。

import numpy

import onnx
from onnx import TensorProto
from onnx import helper, numpy_helper

X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [28, 28])
input_shape = helper.make_tensor_value_info('input_shape', TensorProto.INT64, [1])
reshape_out = helper.make_tensor_value_info('reshape_out', TensorProto.FLOAT, [784])

# numpy array から TensorProtoを作る
input_shape_init = numpy_helper.from_array(numpy.array([784], dtype=numpy.int64))
input_shape_init.name = 'input_shape'

flat_op = helper.make_node(
    'Reshape',
    inputs=['X', 'input_shape'],
    outputs=['reshape_out'],
    name="reshape_node"
)

graph_def = helper.make_graph(
    [flat_op],
    'my-mlp',
    [X],
    [reshape_out],
    initializer = [input_shape_init]
)

model_def = helper.make_model(
    graph_def,
    producer_name='natsutan'
)

onnx.save(model_def, 'onnx/my_mlp.onnx')
onnx.checker.check_model(model_def)
print('The model is checked!')