ONNXの最適化ツールが便利

ONNXで推論を行うと

import onnx
import onnxruntime

onnxfile = "models/hoge.onnx"
session = onnxruntime.InferenceSession(onnxfile)

以下のような警告が出る場合がある

[W:onnxruntime:, graph.cc:3487 CleanUnusedInitializersAndNodeArgs] Removing initializer 'huga'. It is not used by any node and should be removed from the model.

使われなくなったイニシャライザが残ったままになってしまっているという警告なのだが、該当するイニシャライザをいちいち探して削除するのも面倒
何か良い方法は無いのかとググったところ
ONNX Optimizer というライブラリでONNXを最適化する方法が
https://github.com/microsoft/onnxruntime/issues/1899#issuecomment-840596510 にて紹介されていたので
それを少し改変して使用

from pathlib import Path

import onnx
import onnxoptimizer

onnxfile = Path("models/hoge.onnx")
onnx_model = onnx.load(onnxfile)
passes = ["eliminate_unused_initializer"]
optimized_model = onnxoptimizer.optimize(onnx_model, passes)
output_onnx_path = f"{onnxfile.parent}/{onnxfile.stem}_clean.onnx"
onnx.save(optimized_model, output_onnx_path)

不要なイニシャライザを削除したいだけなので今回は passes に eliminate_unused_initializer を指定したが、
他にも任意の最適化オプションを以下から指定可能
https://github.com/onnx/optimizer/tree/master/onnxoptimizer/passes
どういう最適化を行うのかを詳しく知りたい場合は、ヘッダの中に説明が書かれているので、それを読むと良い

https://github.com/PINTO0309/onnx2json を使用して、最適化する前と後のJSONを比較してみると、
不要なイニシャライザが削除されていることが確認できる

参考サイト

• ONNXモデルのチューニングテクニック (基礎編)
• https://github.com/microsoft/onnxruntime/issues/1899#issuecomment-840596510
• https://github.com/onnx/optimizer