ONNXで fuse_consecutive_log_softmax 最適化

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

最適化をしている箇所はここ。log -> softmaxではなく、softmax -> logの並びの時に最適化が有効。

最適化前のグラフ

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

passes = ['fuse_consecutive_log_softmax']

optimized_model = optimizer.optimize(model_def, passes)

最適化後のグラフ

上手くLogSoftmaxに置き換わっています。

全ソース。

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

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

softmax = helper.make_node(
    'Softmax',
    axis = 1,
    inputs = ['X'],
    outputs = ['softmax_out'],
)


log = helper.make_node(
    'Log',
    inputs = ['softmax_out'],
    outputs = ['Y'],
)


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

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

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

onnx.checker.check_model(model_def)

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

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