[論文レビュー] Learning to Optimize Tensor Programs

Tianqi Chen, Lianmin Zheng+6 more|arXiv (Cornell University)|May 21, 2018

Parallel Computing and Optimization Techniques参考文献 36被引用数 124

ひとこと要約

本論文は AutoTVM を提案する。これは学習ベースのフレームワークで、ドメイン特化のコストモデルを構築し、ディープラーニングワークロード向けのテンソルプログラム実装を自動検索・最適化する。

ABSTRACT

We introduce a learning-based framework to optimize tensor programs for deep learning workloads. Efficient implementations of tensor operators, such as matrix multiplication and high dimensional convolution, are key enablers of effective deep learning systems. However, existing systems rely on manually optimized libraries such as cuDNN where only a narrow range of server class GPUs are well-supported. The reliance on hardware-specific operator libraries limits the applicability of high-level graph optimizations and incurs significant engineering costs when deploying to new hardware targets. We use learning to remove this engineering burden. We learn domain-specific statistical cost models to guide the search of tensor operator implementations over billions of possible program variants. We further accelerate the search by effective model transfer across workloads. Experimental results show that our framework delivers performance competitive with state-of-the-art hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPU.

研究の動機と目的

• テンソル演算子の自動最適化を、手作業で調整されたライブラリを超えて追求する。
• ハードウェア実行時間を最小化する大規模なテンソルプログラムスケジュール空間を探索する問題を形式化する。
• 数十億のプログラムバリアントを検索するための学習ベースのコストモデルを開発する。
• 新しいワークロードやハードウェアターゲット間で知識を再利用する転移学習を可能にする。

提案手法

• 与えられたインデックス式 e とコード生成器 g に対してコスト f(x) が未知だが測定可能である、スケジュール探索空間 Se を定義する。
• 低レベルの AST からのドメイン特徴を用いた勾配ブースト木 (GBT) と AST の TreeGRU 埋め込みを用いた2つのコストモデルを開発する。
• 実行時間を予測し選択を導くため、回帰または順位ベースの目的でモデルを訓練する。
• シミュレーテッド・アニーリングと多様性を考慮した選択を組み合わせた探索ループを使用して候補スケジュールを生成し、ハードウェア上での実際の性能を測定する。
• 転移学習を取り入れ、グローバルな不変表現モデルとローカルなドメイン内モデルを組み合わせて新しいワークロードの最適化を高速化する。
• AutoTVM（TVM ベース）が外部の演算子ライブラリなしでバックエンド全体でエンドツーエンドの利得を達成することを報告する。

実験結果

リサーチクエスチョン

• RQ1 学習ベースのコストモデルは、数十億規模のスケジュール空間を効果的に探索して、実際のハードウェア上で高速なテンソルプログラム実装を見つけることができるか。
• RQ2 転移学習は異なるワークロードとハードウェアターゲット間でより高速な最適化を可能にするか。
• RQ3 この領域で高速なスケジュールを見つけるのに最も適した目的は回帰かランキングか。
• RQ4 多様性を考慮した探索と不確実性推定はこの設定で有益か。

主な発見

• 統計的コストモデル（GBTと TreeGRU）は、ハードウェア評価回数を少なくしても、より速いテンソルプログラムを特定するのにブラックボックス系のベンチマークを上回る。
• ランキングベースの目的は、一般に回帰目的よりも、候補プログラムを速度で並べ替える点で同等かそれ以上の性能を示す。
• 多様性を考慮した探索は影響がほとんどないか控えめで、特定のワークロードでいくつかの利得が見られた。
• 転移学習は、以前に見られたワークロードからの知識を GPU およびモバイルデバイス間で再利用することで2倍から10倍のスピードアップをもたらす。
• AutoTVM によるエンドツーエンドの推論性能の改善は、NVIDIA TITAN X、ARM Cortex-A53、ARM Mali バックエンドで 1.2x から 3.8x の範囲。