ONNXで extract_constant_to_initializer 最適化

ONNXがサポートしている最適化 extract_constant_to_initializer を調べてみました。

最適化をしているソースはここ。
https://github.com/onnx/onnx/blob/master/onnx/optimizer/passes/extract_constant_to_initializer.h

ConstantをInitializerに変える最適化ですが、いまいちメリットがわからない。

Addの片方をConstantにして試してみましょう。

最適化前のグラフを見ても、ConstantがnetronによってInitializer扱いされていて最適化後も見た目上は変化なし。

しょうがないので、最適化前後のinitilaizerを表示してみました。

# 最適化パスを指定
passes = ['extract_constant_to_initializer']

optimized_model = optimizer.optimize(model_def, passes)
onnx.save(optimized_model, 'onnx/extract_constant_to_initializer_optimized.onnx')

print("BEFOR:initilaizer", model_def.graph.initializer)
print("AFTER:initilaizer", optimized_model.graph.initializer)

実行すると最適化前はinitilaizerが[]でしたが、最適化後は、initilaizerが増えているのがわかります。やったね。

BEFOR:initilaizer []
AFTER:initilaizer [dims: 1
dims: 2
data_type: 1
float_data: 1.0
float_data: 1.0
name: "const_values"
]

全ソース

import numpy as np
import onnx
from onnx import helper
from onnx import TensorProto
from onnx import optimizer

X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 2])
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 2])
values = np.ones((1, 2), dtype=np.float32)

const = onnx.helper.make_node(
    'Constant',
    inputs=[],
    outputs=['const_values'],
    value=onnx.helper.make_tensor(
        name='const_tensor',
        data_type=onnx.TensorProto.FLOAT,
        dims=values.shape,
        vals=values.flatten().astype(float),
    ),
)

add = helper.make_node(
    'Add',
    ['X', 'const_values'],
    ['Y'],
)

graph_def = helper.make_graph(
    [const, add],
    'test-model',
    [X],
    [Y]
)

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

onnx.save(model_def, 'onnx/extract_constant_to_initializer.onnx')

# 最適化パスを指定
passes = ['extract_constant_to_initializer']

optimized_model = optimizer.optimize(model_def, passes)
onnx.save(optimized_model, 'onnx/extract_constant_to_initializer_optimized.onnx')

print("BEFOR:initilaizer", model_def.graph.initializer)
print("AFTER:initilaizer", optimized_model.graph.initializer)