TVMを使って演算する（Python編）

TensorFlow、Chainer、MXNetなどのディープラーニングのフレームワークによって作成した学習済みのモデルを使って、iPhoneやAndroidスマートフォン、Raspberry PIなどの様々なバックエンドデバイスで推論するために、Amazonは昨年の10月にNNVM/TVMと呼ばれるコンパイラを発表しました。

NNVM/TVMについて

NNVM

NNVMは、フレームワークで作成した学習済みモデルを、共通のモデル（以降、便宜上NNVMモデルと呼びます）に変換するためのモジュールです。
NNVMは各フレームワークで生成されたフォーマットの異なる学習済みモデルを、NNVMモデルに変換します。これにより、各フレームワークで生成する学習済みモデルの差異をNNVMで吸収することができます。

NNVMについては、以下のQiita記事がよくまとまっていますので、こちらを参照ください。

【はじめてのNNVM】https://qiita.com/ashitani/items/e85231297247ec036128

TVM

TVMは、CPU（x86、CPU）やGPU（CUDA、ROCm）などのさまざまなバックエンドで、テンソル演算を高速に実行するための、領域特化言語（DSL）を提供するモジュールです。ハードウェア種別に対応するバックエンドを指定して、ユーザがDSLで記述した演算をコンパイルすることで、ユーザが定義した任意のテンソル演算を各ハードウェアで実行することができます。

NNVMモデルには、推論時に実行するテンソル演算が定義されています。これを入力としてTVMでテンソル演算を定義し、バックエンドを指定して実行することで、冒頭で述べた多様なフレームワークで作成した学習済みモデルを、様々なバックエンドやハードウェアで実行することが可能になります。

例えば、Kerasで作成した学習済みモデルを使って、Kerasが搭載されていない環境（モバイルなどですね）で推論したい場合は、以下の流れに従います。

1. Kerasで学習済みモデルを作成する
2. 推論を実行するデバイスへ、学習済みモデルを配置する
3. NNVMを使って、学習済みモデルをNNVMモデルへ変換する
4. TVMを使って、NNVMモデルで定義された計算を実行する

本記事では、TVMを使ってOpenCLで簡単な計算を試みたので、その結果を紹介します。

TVMのインストール

TVMはPythonのパッケージですが、pipを使ってインストールする方法は用意されていません。現状は、TVMのソースコードをビルドしてインストールする必要があります。TVMをインストールする方法については、以下のチュートリアルを参照してください。

【Installation Guide】http://docs.tvmlang.org/how_to/install.html

TVMを使って演算する

TVMのインストールが完了したら、実際にサンプルコードを使ってTVMを使ってみましょう。

サンプルコード

公式のチュートリアルをそのまま使ってもよいのですが、公式のチュートリアルのコードは細かく分割されているのと、実行結果が少々わかりづらいところがあります。このため、ここで紹介するサンプルコードは、公式のチュートリアルをベースに多少手を加えています。

【Get Started with TVM】http://docs.tvmlang.org/tutorials/get_started.html#sphx-glr-tutorials-get-started-py

なお、各コードで実際に行っている処理をコメントとして追加しました。ただしスケジュールの作成部分については、理解できていません。ただ、このあたりのコードをいろいろ改造してみると、出力される実行コードがそれに応じて変わってくるので、計算を実行するデバイスの特性に応じて最適化したい場合に、ここをチューニングすればよいかと思われます。

import tvm
import numpy as np


def main():
    # 演算の定義
    # ここでは単純に、C[] = A[] + B[]のようなN個の要素をもつ配列の演算を定義します
    n = tvm.var("n")
    A = tvm.placeholder((n,), name="A")
    B = tvm.placeholder((n,), name="B")
    C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")

    # スケジュールの作成（よくわかっていないところです）
    s = tvm.create_schedule(C.op)
    bx, tx = s[C].split(C.op.axis[0], factor=1)
    s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
    s[C].bind(tx, tvm.thread_axis("threadIdx.x"))

    # ターゲットをOpenCLとして、TVMでの実行コードを生成
    fadd_cl = tvm.build(s, [A, B, C], "opencl", name="fadd_cl")

    # OpenCLのカーネルコードを表示
    print("=== OpenCL Kernel Code ===")
    print(fadd_cl.imported_modules[0].get_source())

    # 以下のコメントを外すことで、ホスト側のVMで実行されるコードを出力できます
    # print("=== Host VM (Stack VM) Code ===")
    # print(fadd_cl.get_source())

    # 実行コンテキストをOpenCLに設定
    ctx = tvm.opencl(0)

    # 計算に使用するメモリを割り当て＆値を設定（配列の要素数は16）
    # 入力値はランダムに設定
    n = 16
    ndarr_A = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
    ndarr_B = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
    ndarr_C = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)

    # OpenCLで演算実行
    fadd_cl(ndarr_A, ndarr_B, ndarr_C)

    # tvm.nd.arrayは、インデックスアクセスやforアクセスが未サポートのため、NumPyに戻す
    np_A = ndarr_A.asnumpy()
    np_B = ndarr_B.asnumpy()
    np_C = ndarr_C.asnumpy()

    # 演算結果を表示
    print("=== Result ===")
    for i, (a, b, c) in enumerate(zip(np_A, np_B, np_C)):
        print("[{}] {:.4f} + {:.4f} = {:.4f}".format(i, a, b, c))


if __name__ == "__main__":
    main()

実行結果

サンプルコードの実行結果を以下に示します。

=== OpenCL Kernel Code ===
__kernel void fadd_cl__kernel0(__global float* restrict C, __global float* restrict A, __global float* restrict B) {
  C[((int)get_group_id(0))] = (A[((int)get_group_id(0))] + B[((int)get_group_id(0))]);
}


=== Result ===
[0] 0.7892 + 0.2392 = 1.0285
[1] 0.2157 + 0.1431 = 0.3588
[2] 0.0036 + 0.8343 = 0.8378
[3] 0.5737 + 0.4944 = 1.0681
[4] 0.1683 + 0.3953 = 0.5637
[5] 0.6003 + 0.5182 = 1.1185
[6] 0.5591 + 0.6883 = 1.2474
[7] 0.7983 + 0.7930 = 1.5913
[8] 0.3545 + 0.3017 = 0.6562
[9] 0.4745 + 0.8248 = 1.2993
[10] 0.3873 + 0.4303 = 0.8176
[11] 0.3797 + 0.7478 = 1.1275
[12] 0.8680 + 0.0090 = 0.8770
[13] 0.4732 + 0.5619 = 1.0350
[14] 0.0511 + 0.2702 = 0.3213
[15] 0.4075 + 0.0753 = 0.4828

最初に、TVMで生成したOpenCLのカーネルコードを出力しました。OpenCLのカーネルコードに馴染みのない人でも、このコードから足し算を行うコードだと何となくわかるかと思います。

また、このカーネルコードを実行した結果も出力しました。ランダムの値が設定された16個の要素から構成される配列について、足し算できていることが確認できますね。

おわりに

PythonからTVMを使って、配列の足し算という簡単な演算を行ってみました。
文字列でデバイス名を指定するだけで、CUDAやOpenCLなどの計算を実行するデバイスを変更できるところも興味深いですが、各バックエンド向けに出力されるカーネルコードを確認できるのも面白いです。
今後は、理解できていないスケジュールの部分を改造したり、別の計算をTVMで実行させたりしながらTVMに対する理解を深めていきたいです。