ONNXで eliminate_unused_initializer 最適化

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

最適化をしているソースはこれ。

ソースよりもコメントがわかりやすい。

// Before:
//   A, B, C are in the initializer list
//   D = Add(B, C)
// After:
//   B, C are in the initializer list and A is removed
//   D = Add(B, C)
//

グラフのどこにも使っていない initializer があれば削除されます。

ダミーの initializer を定義してグラフに加えます。

dummy_init = numpy_helper.from_array(numpy.zeros((190, 100), dtype=numpy.int64))
dummy_init.name = 'dummy_init'

最適化前のグラフ。

passにeliminate_unused_initializerを指定して最適化。

passes = ['eliminate_unused_initializer']

optimized_model = optimizer.optimize(model_def, passes)

最適化後のグラフ。

グラフを修正するのではないので、グラフの表示は同じです。モデルのファイルサイズが148 KBから148 バイトへ1/1000になりました。

全ソース。

import numpy
import onnx
from onnx import helper, numpy_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, 4])

dummy_init = numpy_helper.from_array(numpy.zeros((190, 100), dtype=numpy.int64))
dummy_init.name = 'dummy_init'

pad = helper.make_node(
    'Pad',
    ['X'],
    ['Y'],
    mode = 'constant',
    value = 1.5,
    pads = [0, 1, 0, 1]
)

graph_def = helper.make_graph(
    [pad],
    'test-model',
    [X],
    [Y],
    initializer = [dummy_init]
)

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

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

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

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