AWS Inferentia 性能評価

Published in

nttlabs

13 min readJan 24, 2020

こんにちは、NTT の森田です。

昨年の AWS re:Invent 2019 で、AWS Inferentia チップを搭載した Inf1 インスタンスが発表されました。この Inferentia のコンパイラですが、私が携わっている TVM が使われているということで、中身を調べつつ性能評価してみたいと思います。

まずは動かしてみる

Inferentia を試すには AWS Neuron SDK が必要になります。サンプルを動かすだけなら、ここの下にあるチュートリアルに従うだけですので、それほど難しくはないです。TensorFlow, PyTorch, MXNet の 3 種類の学習フレームワークがサポートされていますので、好みのものを選んでセットアップします。チュートリアルに従っていくと、ResNet50 による画像分類の推論が実行できるようになっています。

コンパイルの実行は EC2 インスタンス上でなくても可能ですが、かなり CPU とメモリを必要とするので、ハイスペックなマシンで行いましょう。少し大きなモデルをコンパイルしようとすると、あっという間に 100 G超のメモリを消費するので、パソコン等でコンパイルするのは難しいかもしれません。

コンパイラの中身を見てみる

Inferentia のコンパイラは、AWS Neuron のリポジトリから入手できる Neuron Compiler (neuron-cc) になります。neuron-cc には TVM が使われているということで、少し中を覗いてみました。

neuron-cc のフロントエンドには TVM の Relay が使われています。

$ strings .venv/lib/python3.6/site-packages/neuroncc/driver/jobs/support/Frameworks.cpython-36m-x86_64-linux-gnu.so | grep tvm
tvm.relay.frontend.tensorflow
tvm.relay.frontend.mxnet
tvm.relay.frontend.onnx

内部では Tensorflow、MXNet、ONNX がサポートされているようです。PyTorch は直接はサポートされていませんが、neuron-cc で一旦 PyTorch のオペレータをTensorFlow のオペレータに変換することでコンパイルしています。

Inferentia のコンパイルの流れですが、まずニューラルネットワークのグラフから、Inferentia 上で処理できる部分をサブグラフとして抽出し、それぞれのサブグラフに対して neuron-cc を実行してコンパイルを行なっています。Inferentia が対応していないオペレータは CPU 上で実行されます。

具体的に neuron-cc が対応しているオペレータは Neuron SDK のリリースノートや、neuron-cc コマンドから見られます。現在のサポート状況でも、既存の多くのモデルが動かせそうですが、実際にはちょっとモデルが大きかったり複雑だったりするとコンパイルエラーになったり、メモリを使用し尽くして out of memory になったりと、なかなかコンパイルにはコツがいります。この辺はもう少し品質が良くなってくると嬉しいところです。

私個人的には以前紹介した Image Style Transfer （画風変換）を動かしてみたかったのですが、Instance Normalization が TVM レベルでサポートされていないのでダメでした。まずはコミュニティ版の TVM で、このオペレータをサポートさせないとダメですね。

また、AWS Neuron のリポジトリに入っている Neuron 版 TVM を見てみると、tonga というバックエンドが追加されています。これが Inferentia のようです。私は Inferentia のコンパイラは VTA ベースになると思っていたので、これはちょっと意外でした。

ベンチマーク

それでは Inferentia の性能を見ていきたいと思います。基本性能として積和演算の処理性能を、実際に得られる性能として様々なモデルの推論時間を測定します。

ベンチマークの比較対象として、以下の 3 つのハードウェアでも性能評価を行いました。それぞれ Inferentia と同じように TVM でコンパイルして比較します。

NVIDIA Tesla T4 GPU: 推論用 GPU として比較。AWS EC2 G4 インスタンスを使用。CUDA 10.2 + cuDNN 使用。
NVIDA Tesla V100 GPU (PCIe): GPU で出せる最高性能として比較。CUDA 10.2 + cuDNN 使用。
Intel Xeon Gold 6130 CPU (2 CPU): 合計 32 コア/64 スレッドの CPU で、AVX-512 命令を駆使して演算。

Inferentia を動かす深層学習フレームワークとしては、MXNet を選択します。Tensorflow と PyTorch も利用可能なのですが、

PyTorch は対応しているオペレータがまだ少ない
Tensorflow は Inferentia を動かすと MXNet や PyTorch に比べて、なぜか性能が悪い

ということで却下しました。ただ、MXNet も問題があって

Python でマルチスレッドで性能がスケールしない (Tensorflow と違って GIL の影響を受ける？)
マルチプロセスで動かすと neuron-rtd (Neuron Runtime Daemon) のエラーで落ちる

と、Inferentia のマルチコア（NeuroCores と呼ばれるコアが 4 つ搭載）を有効活用できない感じなのですが、これらの問題はいずれ解決するだろうということで、とりあえず現状での最大性能が見積もれそうな MXNet で評価します。

積和演算

Inferentia は、int8 で 128 TOPS、FP16/BF16 で 64 TOPS のスペックということで、実際にどこまで性能が出るか試してみます。MXNet のサンプルを少し修正して、1024x1024 の行列の乗算を 100 回実行するプログラムにします。

行列のサイズと繰り返し回数は、エラーなしでコンパイルできる、ギリギリの大きさです。

プログラムに関してですが、正方行列なので、計算性能を図るだけなら転置 (transpose_b=True) は不要そうですが、これを入れておかないと Relay のレイヤで transpose オペレータが入ってしまうので、あえてオプションをつけています。また、int8 の方が性能が出そうですが、コンパイルエラーになったので float16 で試しました。

このモデルをコンパイルして Inf1 インスタンス上で動かします。

平均の推論時間は 17.1 ms でした。

1024x1024 の行列同士の掛け算を 100 回実行した結果ですので、計算回数は 1024 * 1024 * 1024 * 2 * 100 になります。なので性能は、計算すると 12.55 TOPS です。1 コアのみを使った性能ですので、4 コア使えるとこの 4 倍の 50.2 TOPS 出せることになります。カタログスペックは 64 TOPS なので、かなりハードウェアの性能を引き出せていることがわかります。

他のハードウェアについても性能を測定してみました。

GPU (T4): 5.98 ms (35.9 TFLOPS)
GPU (V100): 3.75 ms (57.3 TFLOPS)
CPU: 5.55 s (38.7 GFLOPS)

T4 の FP16 演算のカタログスペックは 65 TFLOPS、V100 は 112 TFLOPS なので、理想値の 50 % 強の性能が出せていることになります。NVIDIA GPU の性能に関しては cuDNN や cuBLAS で決まってしまうので、どのフレームワークやコンパイラを使っても似たような結果になります。ただ、TVM では cuDNN や cuBLAS を使わないで Volta の Tensor コアを直接サポートする取り組みも進んでいて、cuDNN より高速になりそうな可能性を秘めているので、もう少し開発が進んだら評価してみたいところです。

CPU に関しては極端に遅いですが、積和演算が CPU に不利なベンチマークなので仕方ないところです。ただ、FP32 にして測定すると Intel MKL が有効にできて、性能は 545.2 GFLOPS まで跳ね上がりました。

画像分類

基本性能は Inferentia に期待できることがわかったので、実際のニューラルネットワークモデルで評価してみましょう。まずはシンプル推論処理である画像分類を試します。サンプルの ResNet-50 よりも層が多い、ResNet-152 (v1) で測定しました。モデルは Gluon のものを利用しました。入力画像サイズは 224x224 で、バッチサイズは 1、計算精度は FP32 です。結果は以下の通りになります。

Inferentia: 9.30 ms
GPU (T4): 12.2 ms
GPU (V100): 9.38 ms
CPU: 21.86 ms

Inferentia は V100 並みの性能が出ていますね。1 コアしか使えていないので、もしも 4 コア使えると理想的には 4 並列でも同性能で処理できることになります。ちなみにバッチサイズ 4 で測定すると、T4 は 26.92 ms、V100 は 13.3 ms の性能でした。バッチサイズ 4 以下であれば Inferentia はかなり有用な感じがします。

物体検出

物体検出のモデルはいろいろ試しましたが、うまくコンパイルが通って正しく動かせるものは、ありものでは見つかりませんでした。そこで MXNet で動くシンプルな Tiny YOLO v2 のモデルを作ったところ、うまく動作しましたのでこれで評価します。入力画像は 416x416 で、計算精度は FP32。なお、Non-Maximum Suppression のレイヤは Inferentia にサポートされないので入れませんでした。結果は以下の通り。

Inferentia: 5.29 ms
GPU (T4): 3.07 ms
GPU (V100): 1.39 ms
CPU: 7.20 ms

今度は GPU に比べて Inferentia が遅いですね。原因は調べられていませんが、もしかしたらモデルがかなり軽量な割に、入出力のサイズが多少大きいところが影響しているのかもしれません。また、CPU もなかなか健闘しています。

意味領域分割

意味領域分割 (semantic segmentation) とは、入力画像からピクセルレベルで物体検出を行うタスクです。

Gluon の fcn_resnet101_ade というモデルを使いました。入力画像は 200x200 （これより大きいとコンパイルできない）で、計算精度は FP32 です。

Inferentia: 227.8 ms
GPU (T4): 24.1 ms
GPU (V100): 10.6 ms
CPU: 48.0 ms

Inferentia が T4 の 10 倍近く遅い結果となってしまいました。流石におかしいので、コンパイルしたグラフを可視化して調べてみます。

_neuron_subgraph_op というのが、Inferentia で実行される、モデルのサブグラフです。Inferentia で処理した後、画像の resize をするオペレータ (_contrib_BilinearResize2D) が２つ、CPU 上で処理されるようになっています。この二つが悪さをしているのかもしれません。画像のリサイズはモデルから外して、再度性能測定してみます。

Inferentia: 13.0 ms
GPU (T4): 22.3 ms
GPU (V100): 10.0 ms
CPU: 44.0 ms

これだと V100 と同程度の性能が出ました。なかなかチューニングが難しいですね。

姿勢推定

入力画像の人物の姿勢を、棒人間のような形で推定するタスクです。物体検出と組み合わせることで、下の画像のようなことができます。

モデルは Gluon の simple_pose_resnet152_v1b を使いました。入力画像は 224x224 で、計算精度は FP32 です。結果は以下の通りです。

Inferentia: 38.7 ms
GPU (T4): 15.5 ms
GPU (V100): 10.6 ms
CPU: 28.7 ms

今回も Inferentia が一番遅いので、コンパイルしたモデルを可視化して見てみます。

_neuron_subgraph_op の箇所が、Inferentia で実行されるサブグラフですが、4 つに分断されてしまっていますね。Deconvolution は Inferentia でまだサポートされていないので、この分断は避けられません。一つ目と四つ目のサブグラフは convolution を含んでいるので Inferentia で実行した方が速そうですが、二つ目と三つ目のサブグラフは batch normalization と ReLU しか含んでいないので、Inferentia へのメモリコピーのオーバヘッドを考えると、CPU 上で実行した方が速そうです。

今回の例だと、以下のようにコンパイル時に指定することで、特定のオペレータを強制的に CPU 上で実行させることができます。

excl_node_names = [‘batchnormcudnnoff0_fwd’,
                   ‘relu0_fwd’,
                   ‘batchnormcudnnoff1_fwd’,
                   ‘relu1_fwd’]
sym, args, aux = mx.contrib.neuron.compile(sym, args, aux, inputs,
                                  excl_node_names=excl_node_names)

こうすると、コンパイル結果のグラフは以下のようになります。

Inferentia で実行されるサブグラフが二つに減って、メモリコピーのオーバヘッドが減りました。これで Inferentia の性能測定をすると、実行時間は 32.0 ms になります。サブグラフ四つの時よりは高速になりましたが、まだ GPU や CPU より遅いですね。さらに高速化するには、deconvolution が Inferentia にサポートされないと難しそうです。

おわりに

Inferentia、色々試してみましたが、現状はまだコンパイルに失敗することが多かったりと扱うのは難しかったです。ただ、うまくいくと Tesla T4 (G4 インスタンス)よりも安い価格で高い性能を叩き出すので、推論環境の選択肢には十分なりえそうです。そのうち SageMaker と連携するらしいので、それまでにはもう少しコンパイラの品質が上がってくることを期待です。

私たちNTTは、オープンソースコミュニティで共に活動する仲間を募集しています。ぜひ弊社ソフトウェアイノベーションセンタ紹介ページ及び、採用情報ページをご覧ください。