TVM v0.5 リリース

TVM の新機能紹介

Kazutaka Morita

Published in

nttlabs

10 min readFeb 22, 2019

先日 2019/2/19 に TVM v0.5 がリリースされました。リリース内容は以下のリンク先にあります。

TVM 0.5 Release Note · Issue #2448 · dmlc/tvm

The release is managed by the TVM PMC, moderators for this release are @ZihengJiang and @yzhliu The community welcomes…

github.com

新機能はものすごく多岐にわたっていますが、ユーザへの影響が特に大きいのは

ニューラルネットワークを記述する新しい言語 (Relay) のサポート
自動チューニング機能 (AutoTVM) のチューニング範囲拡張

あたりです。Relay と AutoTVM については、私の前回の記事もご参照ください。

TVM はものすごい勢いで開発が進んでいるため、ドキュメントやチュートリアルの整備が追いついていないです。そのため、リリースノートに載っている機能であっても、使い方が開発者にしか分からないものがいくつかあります。今回は、そういった隠れた新機能のうち、個人的に注目している量子化とヘテロジニアスruntimeについて解説したいと思います。

量子化

TVM v0.5 から量子化機能がサポートされました。ここでの量子化とは、入力の深層学習モデルに対して、例えば 32 ビット浮動小数点数で表現されているパラメータを、8 ビット整数などのより低い計算精度の表現に変換することです。量子化を行うと推論の認識精度は下がることが多いですが、メモリ使用量や推論性能が向上するので、エッジデバイス等に組み込む際には重要な機能となります。

以下の畳み込み演算のサンプルプログラムを例に量子化機能の使い方を見ていきましょう。

このプログラムを実行すると、以下のように Relay の中間表現が表示されます。

fn (%data: Tensor[(1, 1, 16, 16), float32]) {
  %0 = meta[relay.Constant][0]
  %1 = nn.conv2d(%data, %0)
  %2 = nn.relu(%1)
  %2
}

この nn.conv2d（二次元の畳み込み演算）と nn.relu（活性化関数 ReLU）を量子化してみます。先ほど作成した Relay 関数に対して、以下を適用します。

ここの qconfig で様々なパラメータを指定することで、量子化の動作を制御することができます。現在利用可能な設定は以下です。

量子化後のビットサイズを指定: nbit_input, nbit_weight, nbit_activation（初期値はそれぞれ 8, 8, 32）
量子化後のデータ型を指定: dtype_input, dtype_weight, dtype_activation（初期値はそれぞれ "int8", "int8", "int32"）
量子化で値を変換する際のスケール量を指定: global_scale（初期値は 8）
最初に登場するいくつかの畳み込み演算を、量子化しないようにする設定: skip_k_conv（初期値は 1）
データを丸める際にバイアス値の加算を行うかどうかの設定: round_for_shift（初期値は True）
量子化したデータを元に戻す時に、一旦低ビットに丸めるかどうかの設定: store_lowbit_output（初期値は True）

quantize 関数を適用することで、Relay の中間表現は次のように書き換わります。32 ビット浮動小数点数が -127 から 127 までの 8 ビット整数にスケールされている様子などが見えます。

fn (%data: Tensor[(1, 1, 16, 16), float32])
    -> Tensor[(1, 1, 16, 16), float32] {
  %0 = multiply(%data, 16f)
  %1 = round(%0)
  %2 = clip(%1, a_min=-127, a_max=127)
  %3 = cast(%2, dtype="int8")
  %4 = meta[relay.Constant][0]
  %5 = nn.conv2d(%3, %4, out_dtype="int32")
  %6 = nn.relu(%5)
  %7 = add(%6, 64)
  %8 = right_shift(%7, 7)
  %9 = clip(%8, a_min=-127, a_max=127)
  %10 = cast(%9, dtype="int8")
  %11 = cast(%10, dtype="float32")
  %12 = multiply(%11, 0.0625f)
  %12
}

オブジェクト認識で使われる、Tiny YOLO v2 モデルに対して、TVM の量子化を試してみました。結果は下の画像の通りです。量子化することで認識しているオブジェクトが減ってしまっていますが、量子化しても確かにオブジェクト認識はできていることがわかります。

Tiny YOLO v2 によるオブジェクト認識。左がオリジナルのモデルの処理結果で、右が量子化したモデルの処理結果。

ヘテロジニアス runtime

ヘテロジニアス runtime は、様々なデバイスを搭載したヘテロジニアス環境において、複数のデバイスを用いて推論処理を行う機能です。例えば CPU 以外に AI アクセラレータを利用可能な場合、深層学習モデルのどの処理をAIアクセラレータにオフロードするか、を柔軟に設定することができます。

先ほどの畳み込み処理の例では、ニューラルネットワークの全ての処理が GPU 上で実行されていました。これを（あまり現実的ではないですが）

conv2d は GPU 上で実行
relu は CPU 上で実行

となるように書き換えてみましょう。

書き換わった部分は

11行目：conv2d を GPU 上で実行するように指定
13〜17行目：行ったデバイス指定に従って Relay の抽象構文木を変換
24〜25行目：複数のデバイスを用いて推論処理を実行するように指定

になります。これにより、出力される Relay の中間表現は以下のように変わります。 device_copy が挿入されて、GPU メモリからホストメモリへのデータコピーが発生することがわかります。

fn (%data: Tensor[(1, 1, 16, 16), float32]) {
  %0 = meta[relay.Constant][0] # ty=Tensor[(1, 1, 1, 1), float32]
  %1 = nn.conv2d(%data, %0) # ty=Tensor[(1, 1, 16, 16), float32]
  %2 = device_copy(%1, meta[relay.attrs.DeviceCopyAttrs][0]) # ty=Tensor[(1, 1, 16, 16), float32]
  %3 = nn.relu(%2) # ty=Tensor[(1, 1, 16, 16), float32]
  %3
}

実際に nvprof コマンドを使って、どの処理が GPU 上で実行されているか確認してみましょう。まずは、デバイス指定を行わなかった時の nvprof の出力です。

$ nvprof --csv python conv2d.py 2>&1 | awk -F, '/GPU/ {print $8}'
"fused_nn_conv2d_nn_relu_kernel0"
"[CUDA memcpy HtoD]"

GPU 上で実行されている kernel の名前から、conv2d と relu の両方が GPU 上で実行されていることがわかります（TVM の最適化によって、二つのオペレータは一つの kernel に処理がまとめられています）。次に、ヘテロジニアス runtime 機能を使って、conv2d のみを GPU で実行した場合の nvprof の出力です。

$ nvprof --csv python hetero.py 2>&1 | awk -F, '/GPU/ {print $8}'
"fused_nn_conv2d_kernel0"
"[CUDA memcpy HtoD]"
"[CUDA memcpy DtoH]"

kernel の名前から relu の処理が GPU 上で実行されなくなったことがわかります。

ヘテロジニアス runtime は、どの処理をどのデバイスで実行するか、を手動で指定することで利用可能な機能です。ただ、最適なデバイスの割り当て方は自明ではないことが多く、また、検索するパターンも多いので現在のままでは使いにくいです。そのため、将来的には AutoTVM 機能のように、自動で最適なデバイス割り当てを発見する方向へ TVM が進化するのではないかと予想しています。

NNVM から Relay への変換

ここまで説明した TVM の新機能を使うためには、深層学習モデルが Relay で記述されている必要があります。ただ、これまで NNVM で記述してきたモデルを Relay で書き直すことは大変です。そのため、NNVM から Relay へモデルを変換する機能が用意されています。

対応していないオペレータもあるため、どのようなモデルでも変換できるわけではないのですが、有名なモデルに対してはテストもされていてそれなりに動きます。モデルを変換するときは、まずはこの機能で変換を試してみて、それでも失敗したら自分で Relay で書き直してみるのが良いでしょう。

TVM v0.6 に向けて

今回の v0.5 リリースでは、VTA はあまり大きな進展がありませんでした。VTA の新機能に関しては多くが v0.6 リリースに持ち越されており、v0.6 リリースの目玉になりそうです。v0.6 リリースのロードマップは以下で整理され始めています。

TVM v0.6 Roadmap · Issue #2623 · dmlc/tvm

This roadmap for TVM v0.6. TVM is a community-driven project and we love your feedback and proposals on where we should…