TensorFlowでXLAを使えるようにする

このブログ記事は2017年3月22日に開催されたRecap of TensorFlow DEV SUMMIT 2017の私が担当したXLAに関する発表の補足資料です。
XLAの概要を知りたい方は発表資料を御覧ください。

TensorFlow1.0にXLAというコンパイルの仕組みが実験的に導入されましたので、それを試す方法を記します。
XLAの概要や狙いなどはGoogle Developers Japanの記事が参考になります。

XLAはJITとAOTの2つのコンパイラから使われます。
御存知の通り、JITはランタイム上で実行時にコンパイルされ、AOTはランタイムを含まない静的な実行バイナリ（Cライブラリ）を作るコンパイラです。

ここでは、JITとAOTを試しに使ってみる方法をご紹介します。
まずは「A: TensorFlowのビルド」を説明し、「B: JITコンパイルを試す」でJITコンパイルする方法について説明し、「C: tfcompileを試す」でAOTコンパイラのビルドと使い方について説明します。

A: TensorFlowのビルド

プロダクションではXLAがどうやらONになっていないので、自分でXLAがONのTensorFlowをビルドします。

（※基本的にオフィシャルドキュメントに従えばビルドできるはず…）

0. 環境

GCEのGPUインスタンス（Ubuntu16.10 + vCPU:4 + RAM:15GB + Tesla K80）

1. ビルドツール「bazel」のインストール

オフィシャルサイトの手順に従えば躓くこと無くインストールできます。
私はrecommendedになっている「Using Bazel custom APT repository 」にしたがってインストールしました。

2. TensorFlowの依存ライブラリをインストール

$ sudo apt-get install python-numpy python-dev python-pip python-wheel

3. NVIDIA系ライブラリのインストール

以下の3つをインストールします。インストール方法はコチラを参考にどうぞ。
CUDA Toolkitとそれに準拠するドライバのインストール
ディープラーニング用ライブラリcuDNNのインストール
libcupi-devのインストール

4. TFをビルド

まずソースコードをcloneしてきます。

$ git clone https://github.com/tensorflow/tensorflow

gccのバージョンを4にします。
以下のコマンドでgccのバージョンを確認します。

$ gcc -v
gcc version 6.2.0 20161005

nvcc(Nvidia CUDA Compiler)がgcc4までしか対応していないのでgcc-4.9をインストールします。

$ sudo add-apt-repository ppa:ubuntu-toolchain-r/test
$ sudo apt-get update
$ sudo apt-get install gcc-4.9

コンフィグを設定します。

$ ./configure

以下の設定項目だけ注意して設定して下さい。
・gccのパスは先にインストールしたgcc-4.9を指定
・CUDA supportをON
・XLA just-in-time compilerをON

コンフィグが整ったらビルドします。

$ bazel build --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package

凄く時間がかかります。私の環境では2時間ぐらい掛かりました。
ちなみに、TF自体をCPU環境に特殊化し高速化したい場合、このStack Overflowのようにビルドコマンドにオプションを追加するとSSEやAVX、FMAが利用できます。
ビルドが完了したらwheelパッケージを作ります。

$ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

これで/tmp/tensorflow_pkg配下に「tensorflow-1.0.1-cp27-cp27mu-linux_x86_64.whl」というwheelファイルが作られます。

あとは、これをpip installするだけですが、実験要素を含んだ不安定なTensorFlowなので、virtualenvの中にインストールしておいたほうが良いと思います。

$ virtualenv --system-site-packages .venv
$ source .venv/bin/activate
$ pip install /tmp/tensorflow_pkg/tensorflow-1.0.1-cp27-cp27mu-linux_x86_64.whl

これでXLAが使えるTensorflowのビルドとインストールが完了しました。
では、実際にJITコンパイルしてみます。

B: JITコンパイルを試す

サンプルコードがあるので、そちらを使います。
JITの使用方法をオフィシャルサイトで確認しておくと、理解しやすいと思います。

$ cd tensorflow/tensorflow/examples/tutorials/mnist/
$ python mnist_softmax_xla.py --data_dir=./data

実行が完了するとtimeline.ctf.jsonというファイルが生成されますので、それをChrome　Trace Event Profileで表示できます。
Chromeを開いてURLに「chrome://tracing」と入力するとファイルを読み込めるので、そこからtimeline.ctf.jsonを読み込みます。

残念ながら、現バージョン(1.0.1)のソースコードで私がビルドしたTensorFlowでは、XLAのJITコンパイルが動いていないようでした…（対象のソースコードを色々とイジって試してみたのですが、どれもJITコンパイルされませんでした）
うまくJITコンパイルされると、
これ↓が

こう↓なるみたいです。

_XlaLaunchというOperationがJITコンパイルされたOperationということらしいです。

※ちなみに、1.0.1で公開されているサンプルプログラムのmnist_softmax_xla.pyでは--xlaオプションが効かない（単純にこのプログラムのバグ）ため、XLAをON/OFFしたい場合は、ソースコードを修正する必要があります

B: tfcompileを試す

tfcompileはXLAの仕組みを使ってTensorFlowのコードを実行可能なライブラリにコンパイルしてくれるツールです。
何が凄いって、TensorFlowのランタイムを必要とせずTensorFlowのグラフを実行できるので、実行ファイルのサイズが小さくなったり、レイテンシを低下させたりと、推論フェーズ（実際にMachine Learningをプロダクションで利用するフェーズ）に対して強力なツールとなります。
推論エンジンとして最適化するとGPUがなくてもレイテンシを上げることができるので、サーバではCPUインスタンスで運用できたり（GPUインスタンス高いしドライバとかアップデートの問題とか運用に手間が必要…）、モバイルでの利用も見えていきます。
XLAの本命はこれだと個人的に考えています。

1. tfcompileをビルドする

TFをビルドするときと同じようにbazelを使います。

$ bazel build --config=opt --config=cuda //tensorflow/compiler/aot:tfcompile

これだけでtfcompileが完成します。
実際tfcompileはbazelのマクロとして利用します。

2. サンプルで動作を確認

動作概要をオフィシャルサイトで確認しておくと理解しやすいかと思います。
まずグラフデータが必要になりますので、グラフデータを作ります。

$ cd tensorflow/compiler/aot/tests
$ python make_test_graphs.py

これはいくつかのパターンのグラフデータを作成するプログラムです。
実行するとtest_graph_tfmatmul.pbというファイルが作成されます。
これがグラフデータです。
これにconfigファイルを足すと、実行可能なバイナリとヘッダーファイル（すなわちライブラリ）が生成されます。

以下のコマンドがそれを実行してくれます。

$ bazel build :test_graph_tfmatmul

対応するBUILDの設定はこのようになっています。

tf_library(
    name = "test_graph_tfmatmul",
    testonly = 1,
    config = "test_graph_tfmatmul.config.pbtxt",
    cpp_class = "foo::bar::MatMulComp",
    graph = "test_graph_tfmatmul.pb",
    tags = ["manual"],
)

configでConfigurationファイル”test_graph_tfmatmul.config.pbtxt”を指定しています。
cpp_classでC++のクラス名（名前空間付き）を指定します。
graphは先ほどのグラフファイルを指定します。
Configファイルの中身は以下のようになっています。

feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}
fetch {
  id { node_name: "x_y_prod" }
}

feedがグラフの入力ノードで受け取る多次元配列の形を規定します。ノードが複数あれば、その数分feedを規定します。
fetchは出力ノードです。こちらも同じく、出力ノードの数だけfetchを規定します。

これをビルドすると、ビルド結果が//bazel-genfiles/tensorflow/compiler/aot/testsの中にtest_graph_tfmatmu.{h|o}が生成されます。この場所はbaselのライブラリ/includeパスに含まれているので、生成されたファイルを移動させる必要はありません。

ライブラリが完成しました。
これを利用するプログラムを書きます。こちらもbazelのビルドツールを利用すると以下のようになり簡単です。

cc_binary(
    name = "my_binary",
    testonly = 1,
    srcs = [
        "my_code.cc",
    ],
    deps = [
        ":test_graph_tfmatmul",
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)

nameは生成される実行ファイルの名前を指定します。
depsは依存するライブラリです。「:（コロン）」始まりはbazelのタスク名です。
srcsは実際のソースコードです。先ほど作成したヘッダーファイルをincludeして利用するプログラムです。

my_code.ccの中身は以下のとおりです。

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include 
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());

  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

ライブラリのヘッダーファイルをincludeして、そのクラス名のfoo::bar::MatMulCompを使ったプログラムです。
グラフの出力ノードの値が正しければ（58なら）、「Success」と標準出力するプログラムです。

これをビルドします。

$ bazel build :my_binary

ビルドされた実行ファイル（my_binary）は下記に配置されます。

//bazel-bin/tensorflow/compiler/aot/tests/my_binary

実行するとたしかに「Success」という文字を得ることが出来ます。

C: おわりに

TensorFlow1.0からexperimental機能として実装されたXLAについて、JIT/AOTの両方で使う方法と実際に使う手順を説明しました。
まだうまく動かなかったり、思ったような性能を得られなかったりするので、あくまで実験機能として使うと良いと思います。

JITは学習フェーズの"自動"高速化、AOTは推論フェーズの低レイテンシ・省メモリ化など、Machine Learningの開発から運用までがトータルにサポートされる期待の機能ですので、今後も動向をチェックしていきたいと思います。

adamrocker

2017/03/22