ONNXがサポートしている最適化 fuse_consecutive_squeezes を調べてみました。
最適化をしている箇所はここ。
https://github.com/onnx/onnx/blob/master/onnx/optimizer/passes/fuse_consecutive_squeezes.h
コメントの方がわかりやすい。
// Before:
// X is a tensor with shape=[1, 1, 2, 3, 1, 5, 1]
// Y = Squeeze(X, axes=[1, 4]) -> shape=[1, 2, 3, 5, 1]
// Z = Squeeze(Y, axes=[0, 4]) -> shape=[2, 3, 5]
// After:
// Z = Squeeze(X, axes=[0, 1, 4, 6])
連続するSqueezeがまとめられるなら、一つのSqueezeにfuseします。
最適化前のグラフ
passにfuse_consecutive_squeezesを指定して、optimizer.optimizeを呼び出します。
passes = ['fuse_consecutive_squeezes']
optimized_model = optimizer.optimize(model_def, passes)
最適化後のグラフ
上手く消えてました。
全ソース
import onnx
from onnx import helper
from onnx import TensorProto
from onnx import optimizer
# Before:
# X is a tensor with shape=[1, 1, 2, 3, 1, 5, 1]
# Y = Squeeze(X, axes=[1, 4]) -> shape=[1, 2, 3, 5, 1]
# Z = Squeeze(Y, axes=[0, 4]) -> shape=[2, 3, 5]
# After:
# Z = Squeeze(X, axes=[0, 1, 4, 6])
X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 1, 2, 3, 1, 5, 1])
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [2, 3, 5])
squeeze0 = helper.make_node(
'Squeeze',
axes = [1, 4],
inputs = ['X'],
outputs = ['squeeze0_out'],
)
squeeze1 = helper.make_node(
'Squeeze',
axes = [0, 4],
inputs = ['squeeze0_out'],
outputs = ['Y'],
)
graph_def = helper.make_graph(
[squeeze0, squeeze1],
'test-model',
[X],
[Y]
)
model_def = helper.make_model(
graph_def,
producer_name='onnx_example'
)
onnx.save(model_def, 'onnx/fuse_consecutive_squeezes.onnx')
onnx.checker.check_model(model_def)
# 最適化パスを指定
passes = ['fuse_consecutive_squeezes']
optimized_model = optimizer.optimize(model_def, passes)
onnx.save(optimized_model, 'onnx/fuse_consecutive_squeezes_optimized.onnx')