ONNXの最適化まとめ

ONNXの最適化を一通り試してみたのでまとめ。

サポートしている最適化一覧の取得

サポートしている最適化は、get_available_passesで取得できます。

from onnx import optimizer
all_passes = optimizer.get_available_passes()

大きく分けると、このように分類できます。

• 意味のないOpの削除 （eliminate_deadend等）
• 2つのOpのfusion　（fuse_matmul_add_bias_into_gemm等）
• Convへのfusion　（fuse_add_bias_into_conv等）
• その他

convへのfuseは全く動かず、バージョンアップ待ちです。

最適化の結果

Qiitaにそれぞれまとめました。

• ONNXでeliminate_deadend 最適化
• ONNXで eliminate_identity 最適化
• ONNXで eliminate_nop_dropout 最適化
• ONNXで eliminate_nop_monotone_argmax 最適化
• ONNXで eliminate_nop_pad 最適化
• ONNXで eliminate_nop_transpose 最適化
• ONNXで eliminate_unused_initializer 最適化
• ONNXで extract_constant_to_initializer 最適化
• ONNXで fuse_add_bias_into_conv 最適化（駄目でした）
• ONNXで fuse_bn_into_conv 最適化（駄目でした）
• ONNXで fuse_consecutive_concats 最適化
• ONNXで fuse_consecutive_log_softmax 最適化
• ONNXで fuse_consecutive_reduce_unsqueeze 最適化
• ONNXで fuse_consecutive_squeezes 最適化
• ONNXで fuse_consecutive_transposes 最適化
• ONNXで fuse_matmul_add_bias_into_gemm 最適化
• ONNXで fuse_pad_into_conv 最適化（駄目でした）
• ONNXで fuse_transpose_into_gemm 最適化
• ONNXで split 最適化

nop

nopは最適化のパスとして指定できますが、何も最適化しません。他の最適化のテンプレートです。

lift_lexical_references

これは上手く試すことができませんでした。Loopや条件分岐で階層的になってしまっているグラフをフラットにするようですが、動かすことができませんでした。バージョンアップ時に試してみたいです。

split最適化

これが唯一面白かった最適化です。

グラフの中で毎回計算しないといけない所と、一度だけ計算すればよい所を別のグラフに分離します。

このグラフで考えます。

左上のAddが定数のみの足し算です。グラフの計算を実行する時に毎回計算する必要が無い部分です。真ん中のAddは、入力Xによって結果が変わるので毎回計算必要です。これを、split最適化を行ってグラフを分離すると以下の2つになります。

左が一度だけ計算すればよいグラフ、右が（その結果も使って）推論時に毎回計算するグラフになります。

コンパイラでよくある定数の畳込み（定数伝播）を実現するための最適化です。ONNXの場合実行時の計算精度がわからないのでホスト（最適化の実行環境）で畳込んでしまうのではなく、畳み込みの計算をonnx runtime側に計算させる意図があるのかなと思いました。

最後に

ONNXの勉強も兼ねて順番にやってみましたが途中で飽きてしまいました。最後のsplit最適化が面白くてよかったです。

グラフの最適化はどのフレームワークでもやるので、個別のフレーワークでそれぞれ書かずにONNXでやるのは合理的なきがしましたが、この程度ならそれぞれ書いてもよいですね。操作のためのAPIは、もう一個上のレイヤー欲しい。

Qiitaは書けば書くほど、誰もいいねしていない記事ですら検索の上にくるので、微妙な気持ちになりました。検索ページの3ページ目くらいで出てくれれば良いんですが、そういうのは難しいのかな。今回、Qiitaとはてなに同じような記事を書いたのも、SEO力の比較をしたかったというのがあります。どうなることやら。