ONNX 2020

R&D チームの奥村（@izariuo440）です。相変わらず深層学習モデルの推論に取り組んでいます。深層学習モデルの推論器として二年前に ONNX Runtime を軽くレビューしましたが、当時と比較するとかなり進歩しており、ONNX Runtime だけでなく ONNX 自体や関連ソフトウェアも成熟が進んでいるので、備忘録として私がお世話になっているものをかいつまんで紹介します。

OPTiM TECH BLOG Advent Calendar 2020 12/18 の記事です。

ONNX や ONNX Runtime は二年前の記事で少し解説しています。必要に応じてご参照ください。

tech-blog.optim.co.jp

ONNX

ONNX は Open Neural Network Exchange の略称で、ニューラルネットワークなどの学習済みモデルのための交換可能なフォーマットを提供しています。既に多くの深層学習フレームワークが ONNX 形式のインポート・エクスポートに対応しているので、ある深層学習フレームワークで学習した学習済みモデルを、他の深層学習フレームワークや ONNX Runtime や OpenCV などで推論できます。

本稿執筆時点での最新版は v1.8.0 です。

チュートリアル

ONNX Tutorials には、さまざまな深層学習フレームワークから ONNX モデルをエクスポートするための手順と、エクスポートしたモデルを推論するための手順が Jupyter Notebook で紹介されています。

• エクスポート
  • Caffe、Caffe2、Chainer、CNTK、CoreML、Keras、LibSVM、LightGBM、MATLAB、ML.NET、MXNet、PyTorch、SciKit-Learn、SINGA、TensorFlow
• 推論
  • Caffe2、CNTK、CoreML、MATLAB、Menoh、ML.NET、MXNet、ONNX Runtime、SINGA、TensorFlow、TensorRT、Windows ML
• 量子化
  • onnxruntime/Bert-GLUE_OnnxRuntime_quantization.ipynb at master · microsoft/onnxruntime

最近では ONNX Runtime を使った INT8 量子化にも対応しており、モデルの軽量化や高速化に役立ちます。

エクスポートは、多くの深層学習フレームワークが対応しているものの、Opset バージョンの互換性が不十分な場合がいくつかあり、出力された ONNX モデルに多少手を入れる必要があるのが実情です。いわゆる罠あり状態で、お化粧直しが必要です。

例えば、MXNet からエクスポートされた BatchNormalization ノードの spatial 属性は 1 にしてやる必要があったり、PRelu の slope の shape を [1, C, 1, 1] に修正してやる必要があったり、initializer の型が double になっていることがあり推論器によっては読み込めないので float32 に修正してやる必要があったり、PyTorch が出力した Clip ノードは最大値が指定されていない場合、空文字の initializer を参照しようとするので修正してやる必要があったりします。私のおすすめの Opset バージョンが 11 です。

ONNX Model Zoo

ONNX Model Zoo には、多くの訓練済みモデルが登録されています。以前は画像分類が主でしたが、最近は SSD、Mask R-CNN、YOLOv4 などの物体検出器など、幅広いジャンルのモデルが登録されています。

顔認識関連では、顔認証のための ArcFace、顔検出のための UltraFace （本日マージ）、感情認識のための Emotion FerPlus、年齢性別推定のための Age and Gender Classification が登録されています。

オプティマイザ

オプティマイザがいい感じになってきました。v1.9.0 からは ONNX Optimizer に移動するようです。ドキュメントはまだないようですが、以下のコードからどういうことができるのかを知ることができます。

• https://github.com/onnx/optimizer/blob/master/onnxoptimizer/pass_registry.h#L46

ざっくりいうと、以下のようなことができます。

1. 不要なノードの除去
2. 定数ノードを Initializer に
3. いくつかのノードの合成 (Fuse)

その他

Python API の概要がドキュメント化されました。以前はコードを見ながら四苦八苦していたのと比べると、かなりとっかかりやすくなったと思います。モデルの読み込み・書き出しをはじめ、モデルに含まれるテンソルの内容を numpy でいじる方法や、ヘルパー関数を使って ONNX モデルを Python コードで構築する方法や、ONNX モデルの妥当性をチェックする方法や、ONNX モデルを最適化（レイヤーの合成など）する方法や、shape を推論する方法や、ONNX バージョンを変換する方法や、ONNX モデルからサブグラフを抽出する方法や、入力・出力テンソルの大きさを動的にする方法などを把握することができます。

ONNX 関連のソフトウェア

ONNX Runtime

本稿執筆時点での最新版は v1.6.0 です。

現在も非常に活発に開発が進められています。なんと ONNX v1.2.1 以降のすべてのオペレータへの下位互換性を達成しているようです（すごい）。実行プロバイダとして CPU を選択した場合、解釈できないモデルはほとんどないと思われます。また、ONNX の対応オペレーター拡充に伴い、カスタムオペレータで対応する必要もかなり減ってきたという印象です。

深層学習フレームワーク以外で、同じ推論コードが CPU にも CUDA にもさくっと切り替えて実行できるのにはちょっとした感動があります。

TensorRT の実行プロバイダでは、暗黙的なエンジンキャッシュ作成や、FP16 量子化に対応しています。v1.6.0 から、INT8 量子化にも実験的に対応しています。INT8 量子化にはキャリブレーションが必要で、キャリブレーションには少なくとも 100 〜 1000 枚程度の画像データなどが必要です。これを ONNX Runtime のセッション内部で暗黙的に行うのは難しそうなので、正式対応にはもう少し時間がかかると感じています。とはいえ、ONNX Runtime が推論だけでなく訓練にも対応しはじめていてかつ訓練時間の短縮を達成しているので、何気なく対応してくることを期待しています。

実行プロバイダとして Xilinx Visits-AI への対応も進んでいるのも気になります。上手く行けば、ONNX 形式のモデルがあれば Xilinx FPGA にも対応できるようになるということですね。

onnx-tensorrt

• onnx/onnx-tensorrt: ONNX-TensorRT: TensorRT backend for ONNX

NVIDIA GPU で ONNX 形式のモデルを推論するための変換と実行が可能です。TensorRT に含まれている nvonnxparser の実装がこれです（記憶が曖昧ですが TensorRT 5.0 から OSS になりました）。NonMaxSuppression がサポートされていないので、物体検出をやるときに一手間必要になります。

現在ではかなり多くのオペレータに対応しています。以前はサポートされていないオペレータに対応するためにプラグインを用意する必要がありましたが、現在ではほとんどその必要性がありません。

細かい話ですが、PyTorch 1.6 で opset_vesion 11 を指定して ONNX 形式にエクスポートしたときに Clip オペレータが現れたときは注意が必要です。PyTorch は max が指定されていないときに入力テンソルの max を省略するのではなく、空文字を指定してくるため、nvonnxparser で変換しようとすると、ここで引っかかって変換が失敗します。

さらに細かい話ですが、 Where など Boolean テンソルを扱うモデルの場合、INT8 量子化ができないようです。特にエラーが出力されずに FP32 で量子化されてしまうようなので注意が必要です（私は INT8 なのに FP32 と速度が変わらないことで気が付きました）。

本稿執筆時点での最新版は v7.2.1 です。

Netron

• lutzroeder/netron: Visualizer for neural network, deep learning, and machine learning models

ONNX 形式のモデルを分かりやすく可視化してくれるツールです。ブラウザ版もあります。ONNX に限らず、主要な深層学習フレームワークにはほとんど対応しています。

ONNX Simplifier

• daquexian/onnx-simplifier: Simplify your onnx model

PyTorch で訓練済みモデルを ONNX 形式にエクスポートしたあとに Netron で可視化すると、思ったよりごちゃごちゃした結果になっていることがあります。特に Reshape、Resize、Upsample など テンソルの shape が絡んだオペレータを使うときによくそうなります。

pip で導入して簡単に使えます。

pip install onnx-simplifier
pip -m onnxsim in.onnx out.onnx

実装として、ONNX Runtime を利用して指定した入力テンソルを与えたときの出力テンソルの shape を得ているようです。逆にいると、ONNX Runtime が解釈できないモデルには利用できないでしょう。

本稿執筆時点での最新版は v0.2.19 です。

おわりに

最近の ONNX と ONNX Runtime はかなり使い勝手が向上しており、当たり前にできることがかなり増えています。とはいえ、ONNX エクスポート関連はまだ互換性の問題があるので、Python でお化粧直しユーティリティを作って対処しています。また、エクスポートされた ONNX モデルには前処理（RGB/BGR 変換、NCHW/NHWC 変換、平均引く、標準偏差で割る、など）が含まれていないことが多いので、これらを ONNX モデルに追加する機能も作ったりしています。そのうち公開できるといいなぁ・・・。

ONNX Runtime は Android/iOS にも対応していますが、まだ公式バイナリは提供されていません（ソースからビルドすればいけます）。Android NNAPI への対応もプレビュー状態ですし、この辺りが強化されていくと「とりあえず ONNX Runtime でいいじゃん」となりそうです。また、WebAssembly 対応や WASM SIMD を考えると、もうしばらくは TFLite も併用していく必要があるなぁ、という感じです。

オプティムでは、こうした技術に興味がある・作ってみたい、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp