ぱたへね!

はてなダイアリーはrustの色分けができないのでこっちに来た

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にそれぞれまとめました。

nop

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

lift_lexical_references

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

split最適化

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

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

このグラフで考えます。

f:id:natsutan:20191002095626p:plain

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

f:id:natsutan:20191002095652p:plainf:id:natsutan:20191002095704p:plain

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

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

最後に

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

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

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