Distributed TensorFlowを試してみる


Distributed TensorFlow が[公開](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/distributed_runtime/README.md)されました。[中の人の説明](http://qiita.com/kazunori279/items/981a8a2a44f5d1172856)がわかりやすいですが、TensorFlowの分散環境における並列演算を支援する仕組みです。

Dockerイメージを山ほど作って使うのが想定なのでしょうけれども、低火力コンピューティングの底辺にいる私としては、とりあえず自宅のデスクトップPCのUbuntu14.04(64bit)上でサーバを走らせて、MacBookから叩くところまでやってみます。

自宅の環境はデスクトップPCのほうがCPUがしょぼい上にGPU非搭載なので実用上まったく価値が無いのですが、まあ練習ということで。

## サーバのビルド(Ubuntu14.04)

現時点(2016/2/28)ではソースからビルドする必要があります。[TensorFlow公式](https://www.tensorflow.org/versions/r0.7/get_started/os_setup.html#installing-from-sources)にならって環境構築からビルドまでやってみます。

まずbazelのインストール。

```bash
$ sudo add-apt-repository ppa:webupd8team/java
$ sudo apt-get update
$ sudo apt-get install oracle-java8-installer
$ sudo apt-get install pkg-config zip g++ zlib1g-dev unzip
$ wget https://github.com/bazelbuild/bazel/releases/download/0.2.0/bazel-0.2.0-installer-linux-x86_64.sh
$ chmod +x ./bazel-0.2.0-installer-linux-x86_64.sh
$ ./bazel-0.2.0-installer-linux-x86_64.sh --user
```

~/binにbazelがインストールされるので、パスを通します。次にbazel以外の依存パッケージのインストール。

```bash
$ sudo apt-get install python-numpy swig python-dev
```

TensorFlow本体をgitから落としてサーバをビルドします。

```bash
$ git clone --recurse-submodules https://github.com/tensorflow/tensorflow
$ cd tensorflow
$ ./configure
$ bazel build --jobs 2 -c opt //tensorflow/core/distributed_runtime/rpc:grpc_tensorflow_server
```

"--jobs 2" を忘れるとリソース不足でビルドに失敗しました。

gRPC対応なTensorFlowそのものも必要であればソースからビルドする必要があるそうです。サーバ上には必要ないのですが、動作確認のためにビルドします。Ubuntu14.04(64bit)ではインストラクション通りですね。

```bash
$ sudo pip install wheel
$ bazel build --jobs 2 -c opt //tensorflow/tools/pip_package:build_pip_package
$ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
$ sudo pip install /tmp/tensorflow_pkg/tensorflow-0.7.1-py2-none-any.whl
```

だいぶ時間かかりますね。。まあ非力なPCだからしょうがないですけど。

## 動作確認(同一マシン内)

デスクトップPC上でサーバの起動をして、デスクトップPCのクライアントからサーバを叩いてみます。[チュートリアル](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/distributed_runtime/README.md)まんまです。

```bash:bash@desktop(server)
$ bazel-bin/tensorflow/core/distributed_runtime/rpc/grpc_tensorflow_server --cluster_spec='local|localhost:2222' --job_name=local --task_index=0 &
```

```bash:bash@desktop(client)
$ python
>>> import tensorflow as tf
>>> c = tf.constant("Hello, distributed TensorFlow!")
>>> sess = tf.Session("grpc://localhost:2222")
>>> sess.run(c)
'Hello, distributed TensorFlow!'
```

いけました。

## 動作確認(リモート)

次にノートPCからです。クライアント側もソースから再ビルドが必要のようです。

ビルド手順は省略します。OSXなので[こちら](https://www.tensorflow.org/versions/r0.7/get_started/os_setup.html#installation-for-mac-os-x)にしたがって対応すればよいと思うのですが、私の場合、途中で"-Bsymbolicなどというリンカオプションは知らん！"というエラーが出たので、当該箇所のBUILDファイルからそのオプションを外したら通るようになりました。

さらに、できたwheelをpipでインストールした後にimport tensorflowすると

```
ImportError: No module named core.framework.graph_pb2
```

などと言われます。いろいろ調べるとvirtualenvで回避するのが簡単だというので試したところ、

```bash:bash@note(client)
$ virtualenv -p /usr/local/bin/python distributed_tensorflow
$ . distributed_tensorflow/bin/activate
$ pip install /tmp/tensorflow_pkg/tensorflow-0.7.1-py2-none-any.whl
```

で使えるようになりました。ふう。。

なおサーバのほうもビルドしようと思いましたが、下記のようなエラーでリンクが通りませんでした(追記：2016/3/4時点で最新版をgit pullして試すとビルドが通りました）

```
duplicate symbol __ZNK10tensorflow17BuildGraphOptions11DebugStringEv in:
    bazel-out/local_darwin-opt/bin/tensorflow/core/distributed_runtime/libsimple_graph_execution_state.a(simple_graph_execution_state.o)    bazel-out/local_darwin-opt/bin/tensorflow/core/distributed_runtime/libbuild_graph_options.a(build_graph_options.o)
ld: 1 duplicate symbol for architecture x86_64
```

さて、noteからdesktop（home.local）を呼んでみましょう。

```bash:bash@note(client)
$ python
>>> import tensorflow as tf
>>> c = tf.constant("Hello, distributed TensorFlow!")
>>> sess = tf.Session("grpc://home.local:2222")
>>> sess.run(c)
'Hello, distributed TensorFlow!'
```

動きました。これでネットワーク越しにグラフをやりとりできました。

## 使ってみます

さて、使ってみます。

先日Qiitaに上げた[関数近似で学ぶ chainer とディープラーニング](http://qiita.com/ashitani/items/1dc0a54da218ec224ad8)と同じ処理の[TensorFlow版](http://nbviewer.jupyter.org/github/ashitani/jupyter_examples/blob/master/tensorflow.ipynb)を並列化してみます。もともと大した負荷じゃないので分散する意味はないのですが、まあ練習なので。

クラスタ構成は、パラメータサーバ(ps)が2個、お仕事をするマスタサーバ(master)が1個、という構成とします。

サーバの起動引数は、クラスタ構成を--cluster\_specで、サーバの名前とタスク番号を--job\_name,--task\_indexで与えるようです。

```bash@desktop(server)
grpc_tensorflow_server --cluster_spec='master|localhost:2222,ps|localhost:2223,ps_|localhost:2224' --job_name=master --task_index=0 &
grpc_tensorflow_server --cluster_spec='master|localhost:2222,ps|localhost:2223,ps_|localhost:2224' --job_name=ps --task_index=0 &
grpc_tensorflow_server --cluster_spec='master|localhost:2222,ps|localhost:2223,ps_|localhost:2224' --job_name=ps_ --task_index=0 &
```

ドキュメントやヘルプを見る限りはjob\_name=ps, task\_index=0,1 で区別できるはずなのですが、2224ポートに当ててほしいps1が2223を取りに行こうとして失敗し落ちるので、やむなくps,ps\_と名前を変えました。

さて、これで3つのサーバからなるクラスタが出来上がりました。貧乏構成なので全て同じマシン上で動いていますが、別々のコンテナに割り当てるのも簡単そうです。

パラメータサーバを、重み$W$とバイアス$b$で分離します（意味あるかはともかく）。マスタサーバはセッションとエラーのグラフ演算を受け持ちます。

マシン構成とグラフ中のサーバの分担を図にするとこんな感じです。

![distributed_tensorflow.png](https://qiita-image-store.s3.amazonaws.com/0/30340/2ed7d2a1-2f80-27f8-2821-f37a70cc7eed.png "distributed_tensorflow.png")


クライアント側のコードはこうなります。

```python:python@note(client)
import tensorflow as tf
import numpy as np

def get_batch(n):
    x = np.random.random(n)
    y = np.exp(x)
    return x,y

def leaky_relu(x,alpha=0.2):
    return tf.maximum(alpha*x,x)

x_ = tf.placeholder(tf.float32, shape=[None, 1])
t_ = tf.placeholder(tf.float32, shape=[None, 1])

with tf.device("/job:ps/task:0"):
    W1 = tf.Variable(tf.zeros([1,16]))
    W2 = tf.Variable(tf.zeros([16,32]))
    W3 = tf.Variable(tf.zeros([32,1]))

with tf.device("/job:ps_/task:0"):
    b1 = tf.Variable(tf.zeros([16]))
    b2 = tf.Variable(tf.zeros([32]))
    b3 = tf.Variable(tf.zeros([1]))

with tf.device("/job:master/task:0"):
    h1 = leaky_relu(tf.matmul(x_,W1)+b1)
    h2 = leaky_relu(tf.matmul(h1,W2)+b2)
    y  = leaky_relu(tf.matmul(h2,W3)+b3)
    e  = tf.nn.l2_loss(y-t_)

opt=tf.train.AdamOptimizer()
train_step=opt.minimize(e)

with tf.Session("grpc://home.local:2222") as sess:

    sess.run(tf.initialize_all_variables())
    for i in range(10000):
        x0,t0 = get_batch(100)
        x = x0.astype(np.float32).reshape(100,1)
        t = t0.astype(np.float32).reshape(100,1)

        sess.run(train_step,feed_dict={x_: x, t_:t})

        if i%100==0:
            print "loss,", sess.run(e,feed_dict={x_: x, t_:t})
```


結果は単体で動かすよりも激しく遅くなりました（笑）。ネットワークの帯域幅とマシンそのものが遅いのでまあ当然ですが。

ジョブの名前は区別のためだけで、とくにpsだとパラメータサーバだとかいうことはなさそうですね。同様にセッションを投げる先のサーバも誰でもよさそうです。ただまあ、普通はリソースの種類に応じて割り振るのでしょうから名前で区別できたほうが便利ということでしょう。


##　まとめ

本当はデータ並列化まで試したかったのですが、残念ながらそこまで至りませんでした。複数セッションで、たとえばバッチを小分けにして各サーバに学習させ、各サーバの出力したパラメータの微分を集めて、平均してパラメータ更新する、というようなことをするようですね。

とはいえ、ひとまず感触はつかめました。素晴らしいのはコードがすべてクライアント側だけで完結しているところですかね。クラスタを組むのがちょっと面倒ですが、Kubernetesあたりを使って楽に管理できる仕組みを検討中とのことなので楽しみですね。低火力な私には関係ありませんが（笑）

ビルドにはけっこう時間を取られました。とくにOSX版でハマりました。まあ、そのうち公式Dockerイメージが作られるでしょうし、クライアント側もバイナリリリースに取り込まれるでしょう。

これまではGPUに金をかけるのが主流だった学習の高速化ですが、クラウドコンピューティング環境にどれだけ金を突っ込めるかというステージに移行するわけですね。

私のような趣味でやってる人間にはいずれも手が出ないので安いGPUでも探します。そのうち爆安になったAlteraがたんまり入ったボードだか箱だかが出回る日を楽しみにしつつ。

続編を書きました。

- [Distributed TensorFlow でデータ並列を試してみる](http://qiita.com/ashitani/items/dbe76cb9194d60ead9de)
- [Distributed TensorFlow を Google Cloud Platform で動かしてみる](http://qiita.com/ashitani/items/8b52a6b0ca812712a348)