ONNXで eliminate_nop_transpose 最適化

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

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

static bool is_nop_transpose(const std::vector<int64_t>& perm) {
    for (size_t i = 0; i < perm.size(); i++)
      if (perm[i] != (int)i)
        return false;
    return true;
  }

  bool patternMatchPredicate(Node* node) override {
    return (node->kind() == kTranspose && node->hasAttribute(kperm)) &&
        is_nop_transpose(node->is(kperm));
  }

transpose の転置を決める行列(perm)が、[0, 1, 2, ・・] となる場合は、何も転置しないのでnopと同じ扱いで削除します。

こういう転置のNodeを定義して、最適化をかけてみます。

transpose = helper.make_node(
    'Transpose',
    ['pad_out'],
    ['Y'],
    perm=[0,1]
)

最適化前のグラフがこれ。

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

passes = ['eliminate_nop_transpose']

optimized_model = optimizer.optimize(model_def, passes)

不要な転置が削除されました。

全ソースはこちら。

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

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

transpose = helper.make_node(
    'Transpose',
    ['pad_out'],
    ['Y'],
    perm=[0,1]
)

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

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

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

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

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