TVMでyolo v3を速くしてみる

Published in

nttlabs

11 min readJun 11, 2020

こんにちは。NTTの森永です。

皆さん、TVMはご存知ですか？TVMは、深層学習コンパイラと呼ばれるソフトウェア群の一つで、学習済みの深層学習モデルをエッジデバイスからサーバまで様々な環境に最適化しdeployすることに特化したコンパイラです。昨年Apacheプロジェクトに採用されるなど、注目度が高いOSSで、弊社でも何人かの社員がTVMに貢献しています。

TVMというと、関数型言語に影響された内部実装や独特なコードの書かれ方から難しいイメージが持たれがちですが、使ってみる分には結構簡単です。みなさんにが気軽に試せるよう、本記事は基本的な使い方のチュートリアル的な意味合いで、TVMを使い推論モデルを高速化していく手はずを紹介したいと思います。

今回はGluonCVのYolo v3を高速化してみます。というのも、社内の別チームから、TVMでGluonCVのYolo v3を動かしたらすごく遅かったとの噂を聞き、そんなはずはないだろうと、その再検証も兼ねるためです。評価環境には、Tesla T4を搭載したAWSのg4dn.xlargeインスタンスを用います。

1. 普通に動かしてみる

TVMのビルドは、公式ドキュメントに従い簡単に行えます。今回、build/config.cmakeの設定は(USE_CUDA) (USE_CUDNN) (USE_LLVM)の３つをオンにします。加えて性能分析のためデバッガ機能も使いたいので、(USE_GRAPH_RUNTIME_DEBUG)もオンにします。LLVMはビルド済みの10.0を使用しました。

まず、GluonCVのmodelzooからYolo v3の学習済みモデルを読み込み、TVMのMXNetフロントエンド（relay.frontend.from_mxnet）でTVMが扱う汎用グラフ形式のrelayに変換します。

CUDAをバックエンドに指定してビルドします。この時、TVMのビルド時にUSE_CUDNNをONにしていれば、target = ‘cuda -libs=cudnn’とすることで畳み込み層の実行をcuDNNで行うこともできます。

opt_levelで適用される最適化のパスが変わります。NVIDIAのGPU環境ではカーネルの数が大きく減るOpFusionの影響が特に大きいようです。opt_level=4は動かないことが多いので、オススメは3です。

OPT_PASS_LEVEL = {
    "SimplifyInference": 0,
    "OpFusion": 1,
    "FoldConstant": 2,
    "FoldScaleAxis": 3,
    "AlterOpLayout": 3,
    "CanonicalizeOps": 3,
    "CanonicalizeCast": 3,
    "EliminateCommonSubexpr": 3,
    "CombineParallelConv2D": 4,
    "CombineParallelDense": 4,
    "FastMath": 4
}

2. 性能を分析してみる

ここまでで、TVMを使ってモデルをコンパイルし実行することができました。次に、コンパイルしたモデルを通してTVMの性能を確認してみます。ランタイムのtime_evaluatorを使えば、複数回の平均の実行時間を簡単に測定できます。この機能はwarm upの１回分も考慮し実装されています。

(実行結果) Elapsed average time:13569.84 ms

１回の推論あたり13.5秒と噂通りかなり遅いですね。こんな時はデバッガ機能を使えば、処理時間の内訳を詳細に分析することができます。

デバッグランタイムの利用例

デバッガ機能では推論の実行時に、各処理の実行時間を実行順および実行時間のランキング順に表示してくれます。ランキング順で表示される結果の上位を見てみましょう。

Time(us) Time(%) Shape Inputs Outputs
— — — — — — — — — — — — — — — — — — — — — — — — -fused_vision_non_max_suppression 13513800.0 90.625 (1,322560,6) 2 1
fused_add_sqrt_divide_multiply 1365770.0 9.159 (32,) 2 1

どうやらfused_vision_non_max_supressionに全体の90.625%の時間がかかっているようです。名前からこれはNon-maximum Supression（通称：NMS）の処理を指していると推測できます。NMSは、物体検知モデルの後処理で、推論処理の結果のうち上位をバウンディングボックスに変換するものです。調べたところ、どうやらTVMにはGPUのNMS処理を高速化するためのビルドオプション「USE_THRUST」があるようです。試しにconfig.cmakeで(USE_THRUST)をオンにしてからTVMを再ビルドしコードを実行してみると、確かにNMSレイヤーの遅さは解消されました。

(実行結果) Elapsed average time:39.67 ms

とりあえず、TVMが遅いという噂は設定の問題だったようで一安心です。この結果とcuDNNを比較してみます。今回はMXNetからのcuDNNとTVMでtarget = ‘cuda -libs=cudnn’としたものを測定してみました。

ここまでの結果では、TVMのCUDAバックエンドはcuDNNに劣る結果になってしまいました。ですが、TVMの最適化はここで終わりではありません。というのも、TVMにはAutoTVMという自動チューニング機能があるからです。紹介記事のデータによると、AutoTVMを使えば環境によってはcuDNNを上回る性能を発揮するようですので、次章ではこれを使ってみましょう。

3. 自動チューニングしてみる (AutoTVM)

AutoTVMは専門家がやっていたようなループのアンローリングなどの低レベルなコードの最適化を、機械学習により自動で調整してくれます。技術論文はこちらで読むことができます。

今回は、TVMのCUDAのバックエンドに、AutoTVMを適用していきます。尚、cuDNN+AutoTVMの組み合わせて使うことはできません。AutoTVMの最適化は主にtvm上の低レベルなコードで行われるのですが、cuDNNを使う際は畳み込みのアルゴリズム選択からコード生成まで全てcuDNNに委ねる実装になっているためです。

基本的な使い方はAutoTVMのtutorialに従っていきます。以下のようにAutoTVMの設定をtuning_option内に書いていきます。

AutoTVMの設定項目

記述する設定を個別に見ていきましょう。

・log_file_name：チューニングにより発見したスケジュール設定を保存するファイル名です。

・tuner：パラメータの探索アルゴリズムを指定します。xgb (XGBoost), ga (遺伝的アルゴリズム）, grid（グリッドサーチ）, randomの４種類のアルゴリズムが選べます。論文の結果によると、XGBoostが優れているそうです。

・n_trial：各レイヤーにおけるチューニングの試行回数です。多いほど性能が期待できますがその分時間がかかります。

・early_stopping：早い段階で良い設定を見つけた場合にチューニングを切り上げるための数値です。直近のこの数値の範囲の試行回数でより良い設定が見つからない場合にチューニングを完了します。

・measure_option：チューニング中のビルド環境（builder）や測定環境(runner)を指定します。RPCRunnerでは、分散環境でも効率よくチューニングすることができます。

・use_transfer_learning：Trueの場合、異なるタスク間で転移学習します。収束が早めることが期待されます。

この設定とモデルをAutoTVMに与えて、チューニングを実行してみます。

コードからわかるように、tune_tasks自体はAutoTVMのインターフェイスではないので今回のtuning_optionのような書き方は必須ではないのですが、項目が整理しやすいのとtune_tasksの実装が優秀なので、チュートリアル同様に使うのがオススメです。


[Task  1/23]Current/Best: 19.67/2888.14 GFLOPS | Progress: (20/500)

チューニング中は、上記のように進行状況がコンソールに出力されます。畳み込み1レイヤーあたり40~50分程度のチューニングが計15~16時間ほど行われたのち、チューニング結果が保存されたyolov3.logが生成されました。これをビルド時に適用して推論してみましょう。

チューニング結果を適用してビルドする

(実行結果) Elapsed average time: 28.45 ms

確かに1.4倍ほど早くなりましたね。このように、TVMの真価はAutoTVMによって発揮されます。ですが、今回のチューニングでは、まだ一回りほどTVM+cuDNNより5msほど遅いようです。

実はCUDAバックエンドについて、最近のパッチでTensorCore対応が行われています。この機能、現状はバッチサイズ1では使用することができず、今回の評価では使えていないのですが、将来的に対応範囲が広がることで、AutoTVM+TensorCoreでcuDNNにも負けない高速化が期待できそうです。次回はこの機能についても続編として書いてみたいです。

4. 量子化してみる

量子化（Quantize）とは、深層学習モデルのパラメータや計算のbit精度を落とし、モデル自体や実行速度を軽くする手法です。深層学習では、計算精度を落としても、推論の予測精度はそれほど落ちないため、量子化は広く使われています。

TVMにも量子化の機能があるので試してみましょう。これまでの実行では、計算精度が32bitの浮動小数点でしたので、8bitの整数型で実行するようにしてみましょう。TVMの量子化は非常にシンプルに使えます。

TVMでのモデルの量子化

今回は割愛しますが、量子化の設定についてはqconfigのパラメータで細かく設定できます。

（実行結果） Elapsed average time: 4064.97 ms

量子化時のスケジュールの最適化が進んでいないのか、あるいは実行環境やモデルとの相性が悪いのか、量子化前より大分遅いですね。デバッガで確認したところ各convが均等に遅かったため、前章同様にAutoTVMを適用してみました。

（実行結果） Elapsed average time: 10.44 ms

（7/16 追記）：量子化でのチューニング中のエラーが原因で一部レイヤーの実行時間が遅かったため、各タスクのエラーが出なくなるまで再tuneし、結果を更新しました。41.31 ms→10.44 msとなっています。

AutoTVM+量子化のパワーにより、yolo v3を10ms程度で動かすことができました🚀

5. 終わりに

yolo v3のコンパイルを通して、TVMのデバッガやAutoTVM、量子化といった機能を紹介しました。興味を持ったら、ぜひみなさんもお好みのモデルや環境でTVMを触ってみてください。実際に動かすことで、より理解が深まりTVMを好きになれると思います。もしかしたらバグやドキュメントのミス等を発見してコントリビューションにも繋がるかもしれません。

NTTでは、共にOSSのコミュニティ活動を行う仲間を募集しています。興味のある方は、新卒中途問わず、ぜひ採用情報をチェックしてください！