ONNXで fuse_consecutive_transposes 最適化

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

字面から連続する転置を合成してくるはず。

わかりやすい転置を２つ並べてみます。

transpose0 = helper.make_node(
    'Transpose',
    perm = [0, 2, 1, 3],
    inputs = ['X'],
    outputs = ['transpose0_out'],
)

transpose1 = helper.make_node(
    'Transpose',
    perm = [3, 1, 2, 0],
    inputs = ['transpose0_out'],
    outputs = ['Y'],
)

passにfuse_consecutive_transposesを指定して、optimizer.optimizeを呼び出します。

passes = ['fuse_consecutive_transposes']

optimized_model = optimizer.optimize(model_def, passes)

グラフを確認すると、無事fuseされています。

最適化後のtransposeを調べてみると、転置の次元もちゃんとfuseされています。

attribute {
  name: "perm"
  ints: 3
  ints: 2
  ints: 1
  ints: 0
  type: INTS
}

全ソース

import onnx
from onnx import helper
from onnx import TensorProto
from onnx import optimizer

X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 1, 2, 3])
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 2, 1, 1])

transpose0 = helper.make_node(
    'Transpose',
    perm = [0, 2 , 1, 3],
    inputs = ['X'],
    outputs = ['transpose0_out'],
)

transpose1 = helper.make_node(
    'Transpose',
    perm = [3, 1 , 2, 0],
    inputs = ['transpose0_out'],
    outputs = ['Y'],
)

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

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

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

onnx.checker.check_model(model_def)

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

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

print(optimized_model.graph.node)