2017/03/22

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の開発から運用までがトータルにサポートされる期待の機能ですので、今後も動向をチェックしていきたいと思います。

2017/03/14

GCEのGPUインスタンスでTensorFlowを動かす方法

Google Compute Engine(GCE)にGPUインスタンス(インスタンスにGPUを追加できる)が利用できるようになったので、TensorFlow(TF)を動かした。 

1. GPUインスタンスを作る

オフィシャルドキュメントに沿って作るもよし、Cloud Consoleから作るもよし。

GPUが使えるゾーンには現時点(2017/03/14)では制限があるので気を付ける。
OSの選択は気を付ける。GPU(NVIDIAのCUDA)がサポートしているOSにする。
また、各OS毎にCUDAツールキットのインストール方法が異なるので、気を付ける。
私はUbuntu 16.10のOSを選択。

1.1 CUDA Toolkit 8.0をインストールする

インスタンスが生成されたらSSHとかでログインして下記のスクリプトを実行する。
これでCUDA Toolkit 8.0がインストールされる。
#!/bin/bash
echo "Checking for CUDA and installing."
# Check for CUDA and try to install.
if ! dpkg-query -W cuda; then
  # The 16.04 installer works with 16.10.
  curl -O http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-repo-ubuntu1604_8.0.61-1_amd64.deb
  dpkg -i ./cuda-repo-ubuntu1604_8.0.61-1_amd64.deb
  apt-get update
  apt-get install cuda -y
fi

他(バージョンも含めて)のOSの場合、インストールスクリプトが異なるので、オフィシャルドキュメントを参考にして下さい。

2. GPUの動作確認

nvidia-smiを使ってGPUが認識しているかを確認する。
$ nvidia-smi
Tue Mar 14 01:56:30 2017       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 375.26                 Driver Version: 375.26                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla K80           Off  | 0000:00:04.0     Off |                    0 |
| N/A   35C    P8    27W / 149W |     13MiB / 11439MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|    0      2145    G   /usr/lib/xorg/Xorg                              13MiB |
+-----------------------------------------------------------------------------+

コノような内容がコンソールに表示されたらGPU(0番デバイスにTesla K80)が認識されている。

3. cuDNNを追加

すでにドライバとCUDA Toolkitがインストールされているので次に、TensorFlowを動くようにする。
cuDNNはNVIDIAから取得できるが、アカウント登録が必要になるため、コチラから登録する。連絡がくるまで1日要することもあるので、お早めに。
アカウントがあれば先のURLのDownloadボタンから、必要なcudnnライブラリをダウンロードする。
今回の例だと、CUDA 8.0とLinuxの条件に当てはまる「cuDNN v5.1 Library for Linux」が適切。

.tgz形式のファイルがダウンロードできるので、それを展開して、適切な場所にコピーしてインストール完了。
$ tar zxvf cudnn-8.0-linux-x64-v5.1.tgz
$ sudo cp ./cuda/include/cudnn.h /usr/local/cuda/include
$ sudo cp ./cuda/lib64/libcudnn* /usr/local/cuda/lib64 
sudo chmod a+r /usr/local/cuda/lib64/libcudnn*

4. CUDA Compute Capabilityをインストール

$ sudo apt-get install libcupti-dev


5. TFをインストール

TFのインストールも簡単になったものです。下記コマンド1つでインストールできます。
Python2.7とPython3系ではインストール方法が異なるので気をつけて下さい。
TFはまだバージョン更新が激しいのでvirtualenvを作っておくと幸せになれると思います。
$ mkdir tf && cd tf
$ virtualenv --system-site-packages .venv
$ source .venv/bin/activate
(.venv)$ pip install --upgrade tensorflow-gpu

もし、CPUバージョンのTFをインストールしたいならtensorflow-gpuをtensorflowにして下さい。
Python3系を利用している場合はpipをpip3にして下さい。
これでインストールが完了しました。

5.1 動作確認

$ python
>>> import tensorflow as tf
>>> sess = tf.Session()

これでErrorが表示されなければ(Warningは表示されます)正常動作しています。