5月に TVM のオートチューニングフレームワーク、AutoTVM が発表され、その GPU 対応がついに master ブランチにマージされました。論文では cuDNN よりも速い結果が出ており、本当にそこまでパフォーマンスを出せるのか気になっていたため、早速試してみました。本記事では、CUDA と AMDGPU それぞれについて、AutoTVM で学習したスケジュールが、 cuDNN や MIOpen などのベンダーライブラリと比べてどれほどの性能を出せるのか検証します。
CUDA
ここでは、ImageNet のモデルについて、AutoTVM と cuDNN それぞれを使った場合の実行時間を比較します。TVM から cuDNN を呼ぶことで、TVM の Convolution オペレータを cuDNN のものに置き換えることができます。詳しくは、公式サイトにチュートリアルがありますので、そちらをご覧ください。以下の結果では TVM の全 Convolution オペレータを cuDNN のものに置き換えたバージョンを、TVM + cuDNN としています。また、参考までに cuDNN を有効にした MXNet と PyTorch の結果も載せておきます。
ベンチマークに使った GPU は GTX 1070 ti です。結果は以下のようになりました。
単位は全てミリ秒です。
| VGG16 | Resnet50 | Densenet121 | |
|---|---|---|---|
| AutoTVM | 6.68 | 4.03 | 4.64 |
| TVM + cuDNN | 7.10 | 5.14 | 7.14 |
| MXNet + cuDNN | 8.07 | 5.89 | 8.39 |
| PyTorch + cuDNN | 7.79 | 6.55 | 12.3 |
AutoTVM の結果が cuDNN を使ったどの場合よりも速い、という結果になりました。ただし、TVM は上の実行時間を達成するために cuDNN にはできない以下の最適化をしています。
- 各ネットワーク、各レイヤーごとに事前に長時間かけてチューニングしている
- Winograd Convolution のフィルター変換をコンパイル時に行う、Batch norm のスケーリングをコンパイル時にフィルターの重みに埋め込むなどの、Ahead of Time コンパイラだからこそできる最適化をしている
- Convolution + Batch norm + ReLU などの、Convolution とそれに続く要素ごとのオペレータを一つのカーネルにフューズしている
そのため、この結果から "TVM は cuDNN より速い" とは一概には言えず、ネットワーク全体の推論時間が cuDNN を使った場合よりも速い、ということになります。このベンチマークは公平とは言えませんが、デプロイに特化させれば cuDNN 以上のパフォーマンスが出せる、という点は面白いのではないでしょうか。
VGG のように全結合層への入力が大きいネットワークでは、TVM の全結合層に相当するオペレータはかなり遅いです。 上の TVM の VGG 16 の実行時間 6.68 ミリ秒のうち、2.3 ミリ秒が最後の3つの全結合層に費やされています。これを高速化できれば、さらに MXNet や PyTorch との差を広げられます。
TVM, MXNet, PyTorch がインストールされていれば、このスクリプト で上と同様のベンチマークをすることができます。
AMDGPU (ROCm)
TVMは、AMDGPU 向けには AMD 版の cuDNN に相当する MIOpen をサポートしているので、AutoTVM で学習したスケジュールと MIOpen との比較をすることができます。Tensorflow は最近 ROCm でも動くようですが、面倒なので試していません。
ベンチマークに使った GPU は Radeon R9 Nano です。結果は以下になります。
| VGG16 | Resnet50 | |
|---|---|---|
| AutoTVM | 7.44 | 6.45 |
| TVM + MIOpen | 7.18 | 6.13 |
MIOpen にはわずかに届きませんでしたが、MIOpen は cuDNN と同様に Winograd などのカーネルをアセンブリで実装していることを考えると、AutoTVM の結果はかなり良いと考えられます。また、R9 Nano と GTX 1070 ti の理論性能はほぼ同等ですが、CUDA 版の結果と比較しても特にVGG については大きな遜色はなく、AMDGPU でも十分高速に推論ができることがわかると思います。Resnet については CUDA 版と比べて少し開きがあるので要調査です。AMD 版が少し遅いことの理由の一つとして、コンパイラの成熟性が挙げられます。NVCC がすでに 10年以上使われているのに対して、LLVM の AMDGPU バックエンドはまだ比較的新しいです。今後のパフォーマンス改善には期待できると思います。
終わりに
AutoTVM の成果物として、今のところ CUDA 向けには GTX 1080 ti, Titan X (Maxwell)、AMDGPU 向けには Vega Frontier Edition でチューニングしたパラメータが TVM から提供されています。自分のデバイス用に一からチューニングすることも可能ですが、かなりの時間がかかります。自分の環境では、6コア12スレッドのCPU で、公式チュートリアル の通りに VGG16 の全レイヤーをチューニングするのに 8 時間かかりました。ここは、 XGBoost の GPU 版や、XGBoost よりも高速と言われている Light GBM を使えば速くなるかもしれません。また、GTX 1070 ti 向けに自分でチューニングしても、すでに提供されている 1080 ti 用のパラメータを使った場合とで実行時間の差はごくわずかでした。そのため、最初は TVM が提供しているパラメータを使うのがいいと思います。
一年前に TVM の開発者から cuDNN に勝つのが目標と聞いた時には、そんなことが可能なのか疑問でしたが、本当に達成してしまったのはただ凄いとしか言いようがありません。 本記事では TensorRT には触れませんでしたが、GPU 対応が追加されたPR では TensorRT の結果とも比較されています。TensorRT とも互角のパフォーマンスを出せているようですね。
最近は、 TVM 以外にも学習したモデルのデプロイ用フレームワークが増えてきましたが、AutoTVM の登場によってパフォーマンス面での TVM の優位性はさらに高まったように思います。TVM には、パフォーマンス以外にも
- AMDGPU, モバイル GPU を含む、多様なバックエンドのサポート
- コンパイル時のグラフ最適化
- ONNX を含む、様々な Deep Learning フレームワークからのモデルインポート
など、ユニークな点がたくさんあります。本記事を通して、TVM に興味をもつ人が増えてくれれば幸いです。ぜひ試してみてください。