Microsoft の ONNX Runtime を速攻レビュー

Deep Learning / Machine Learning ONNX
このエントリーをはてなブックマークに追加 follow us in feedly

はじめに

オプティムの奥村です。Microsoft が 2018/12/04 に ONNX Runtime を MIT ライセンスでオープンソースとして公開しました。

azure.microsoft.com

ONNX Runtime は 2018/10/16 に Preview として公開されて気になっていましたが、コードが公開されたのでざっと目を通してみて、ONNX Model Zoo に登録されている物体検出モデルの Tiny YOLOv2 を動かしてみました。

ONNX Runtime を大雑把にいうと以下のようになると思います。

  • Menoh、 onnx-tensorrt のような ONNX の推論エンジンのひとつ
  • CUDA・MKL-DNN の計算支援を受けた高速な推論が可能で、nGraph や TensorRT に対応予定
  • Linux/Mac/Windows で動作
    • CPU は ARM もターゲットに入っている
  • API は Python メインで、C# や C/C++ は強化予定
  • 実行エンジンの追加やカスタムオペレータ対応ができる設計

多様なプラットフォーム・計算資源に対応した推論エンジンの需要に対する回答のひとつなのかもしれません。ちなみに、Initial commit は 2018/11/20 で 12 万行を超えるコミットがありました*1。

ONNX Runtime について

ONNX Runtime は、ニューラルネットワークなどの学習済みモデルのデファクトスタンダードになりつつある ONNX 形式の実行エンジンです。学習機能はなく推論特化型です。

ONNX について

ONNX は Open Neural Network Exchange の略称で、ニューラルネットワークなどの学習済みモデルのための交換可能なフォーマットを提供しています。フォーマットは Protocol Buffers で定義されており*2、計算グラフを表現可能です。既に多くの深層学習フレームワークが ONNX 形式のインポート・エクスポートに対応しているので、ある深層学習フレームワークで学習した学習済みモデルを他の深層学習フレームワークで推論したり、Menoh や onnx-tensorrt など ONNX の推論特化型エンジンで推論できます。

ニューラルネットワークの計算グラフで多用される Convolution、Batch Normalization、ReLU といったオペレーター（レイヤー）は Protocol Buffers では定義されておらず、（基本的には）ONNX の定義に手を入れずに拡張可能な構造となっていて、私はこの構造が気に入っています*3。

ONNX では、ニューラルネットワーク向けに以下のようなオペレーターが定義されています。

ニューラルネットワーク以外にも、線形分類、線形回帰、SVM 分類、SVM 回帰といった機械学習にも対応していて、以下のようなオペレーターが定義されています。

ONNX Runtime の目標

ONNX Runtime は、

  1. どんな ONNX モデルでも
  2. 高性能に
  3. クロスプラットフォームで

実行できることを目標に開発されていくようです。ONNX 標準との互換性維持にも努めていくようです。

ONNX Runtime の性能について

ONNX Runtime は CPU でも GPU での実行可能で、別の実行プロバイダをプラグイン提供することもできるようです*4。ざっと見たところ、C++ で IExecutionProvider というインターフェースを実装することで提供できるようです。現在は、CUDA、MKL-DNN による計算支援を受けられるようで、nGraph や TensorRT にも対応していくようです。

クロスプラットフォームについて

2018/12/05 時点で ONNX Runtime は、

  • Linux、Windows、Mac で動作し、
  • Python、C#、C API を提供*5

しています。API とオフィシャルビルドについては以下で参照できます。

Python は Mac x GPU 以外すべて、C# は Windows x CPU のみがオフィシャルビルドが提供されています。Python なら、CPU 向けは pip install onnxruntime、 GPU 向けは pip install onnxruntime-gpu でインストールできるようです。

ONNX Runtime の拡張

実行プロバイダをプラグイン提供できたり、カスタムオペレータを追加する方法 が提供できたりと、拡張可能な設計になっているようです。この手のものは、カスタムオペレータを受け入れることができないと実利用が困難なケースもあるので、素晴らしいですね。

ビルトインされている実行プロバイダは以下のディレクトリにソースコードが配置されています。

cpu、cuda、mkldnn という 3 つのディレクトリがあり、それぞれの実装があります。

CUDA の実行プロバイダについて

CUDA の実行プロバイダの場合、cuDNN が多用されているようです。Convolution なら cudnnConvolutionForward()、ConvTranspose なら cudnnConvolutionBackwardData()、Batch Normalization なら cudnnBatchNormalizationForwardInference() を呼び出しています。前回の記事で「推論特化型では、深層学習モデルに含まれるレイヤーの合成を行うなどして速度向上を図っています」と書きましたが、上記のような API でオペレーター（レイヤー）を実装してしまうと、レイヤーの合成ができなくなると思います。ソースコード全体を把握しているわけではないので、間違えていたらすみません。計算グラフの最適化については今後の対応に注目したいところです。

普通に ImageScaler オペレーターが実装されているあたりは素晴らしいですね。

YOLOv2 を実行する

ImageSclaer オペレーターが実装されているので、それを使った物体検出モデルである Tiny YOLOv2 を実行してみましょう。

まず ONNX Runtime の環境を構築します。今回は Miniconda 環境でやります。画像の読み込みのために OpenCV もインストールします。

conda create -n onnxruntime python=3.6
conda activate onnxruntime
conda install -c conda-forge opencv
pip install onnxruntime

次に Tiny YOLOv2 のモデルをダウンロードして展開します。

wget https://onnxzoo.blob.core.windows.net/models/opset_8/tiny_yolov2/tiny_yolov2.tar.gz
tar xvf tiny_yolov2.tar.gz

ここからは Python です。まずは画像を読み込んで・・・

import cv2
import numpy as np

# YOLO でよく見かける犬の画像を読み込む
img = cv2.imread('dog.jpg')
# Tiny YOLOv2 が受け付ける解像度にリサイズ
img = cv2.resize(img, (416, 416))
# HWC を NCHW に変換
img = np.expand_dims(img.transpose((2, 0, 1)), 0).astype(np.float32)

ONNX Runtime にモデルを読み込ませて、実行すると・・・

import onnxruntime

# Tiny YOLOv2 の学習済みモデルからセッションを作る
session = onnxruntime.InferenceSession('tiny_yolov2/model.onnx')
# 入力名と出力名を取得する
input_name = session.get_inputs()[0].name # 'image' が得られるはず
output_name = session.get_outputs()[0].name # 'grid' が得られるはず

# 推論を実行する
output = session.run([output_name], {input_name: img})[0]

output には、期待通り (1, 125, 13, 13) という shape を持つ出力が得られました。物体検出として完成させるには、この後に信頼度を sigmoid() にかませて、低信頼度のものを省き、信頼度を softmax() にかませて物体のクラスとスコアにして、座標をデコードして・・・と少し面倒な手順を踏む必要があるのでここでは省きます。こういった後処理も ONNX にオペレーターとして標準化されてくるとよいのですが・・・。なんにせよ、簡単に ONNX の学習済みモデルが実行できるのはいいですね。

まとめ

ONNX Runtime をざっとレビューして、物体検出モデルのひとつである Tiny YOLOv2 を実行してみました。深層学習フレームワークなしで CPU・CUDA・MKL-DNN で手軽に学習済みモデルを利用できるのは素晴らしいですね。Menoh は MKL-DNN、onnx-tensorrt は TensorRT のみだったので、TensorRT のサポートが入るとかなりいろんな環境で高速に実行できる環境が手軽に利用できることになります。また、C/C++ API が整備されると、プロダクション環境でもさらに利用しやすくなると思います。

ONNX Runtime High Level Design を見ると、グラフの最適化による高速な実行にも力を入れていくようです。コードには、前回の記事で紹介した TVM も潜んでます。多様なプラットフォーム・計算資源に対応した推論エンジンは需要が高いので、今後も目が離せませんね。

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

www.optim.co.jp

*1:Initial bootstrap commit. · Microsoft/onnxruntime@89618e8

*2:onnx/onnx.proto at master · onnx/onnx

*3:PyTorch には拡張チュートリアルも用意されています

*4:onnxruntime/AddingExecutionProvider.md at master · Microsoft/onnxruntime

*5:C API は Experimental