ONNXで eliminate_identity 最適化

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

ソースで検索するとここが見つかりました。

  bool patternMatchPredicate(Node* node) override {
    return node->kind() == kIdentity;
  }
  bool runTransform(Node* node, Graph&, NodeDestroyType& destroy_current)
      override {
    node->output()->replaceAllUsesWith(node->input());
    destroy_current = NodeDestroyType::DestroyOne;
    return true;
  }

nodeのkindがkIdentityだと、そのノードの入力を出力に置き換えて、ノード自体を削除してるようです。

そもそもの identity の仕様はここ。

入力をそのまま出力するノードです。多分単位行列的な意味で、identity という名前なんでしょう。

実際にやってみましょう。

Identityのnodeが入ったグラフを作ります。

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

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

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

予想通り、Identity のノードが削除されました。

全ソースはこちら

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, 4])

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

I_op = helper.make_node(
    'Identity',
    inputs = ['pad_out'],
    outputs = ['Y']
)


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

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

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

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

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