ChainerのコードをGPU対応にする

この記事は[ACCESS Advent Calendar 2016](http://qiita.com/advent-calendar/2016/access) 23日目の記事です。
ですが、ACCESSとの関係は特になく、私の趣味の範囲の話です :)

Chainerで、コードをざっと書いた後、GPU対応にしたい、ということは良くありますよね（多分）
そこで、GPU対応していないコードをGPU対応する方法を説明します。

知識がほぼ無い状況から書き始めたので、わたしの作業メモになっていますが、何か参考までに。

とりあえず cuda.cupy
================

元のコードが以下のようにあるとします。（このサンプルコードは、[はじめての深層学習プログラミング](https://www.amazon.co.jp/dp/4774185345) からの引用です。とても分かりやすい本なのでオススメです！）

```python
import numpy as np
import chainer.functions as F
import chainer.links as L
from chainer import Variable, optimizers, Chain


class Model(Chain):
    def __init__(self):
        super(Model, self).__init__(
            l1=L.Linear(2, 1),
        )

    def __call__(self, x):
        h = F.sigmoid(self.l1(x))
        return h


model = Model()
optimizer = optimizers.MomentumSGD(lr=0.01, momentum=0.9)
optimizer.setup(model)

x = Variable(np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32))
t = Variable(np.array([[0], [1], [1], [1]], dtype=np.float32))

for i in range(0, 3000):
    optimizer.zero_grads()
    y = model(x)
    loss = F.mean_squared_error(y, t)
    loss.backward()
    optimizer.update()

    print("loss: ", loss.data)

print(y.data)

```

これを、GPU対応してみます。


```python
import chainer.functions as F
import chainer.links as L
from chainer import Variable, optimizers, Chain, cuda


class Model(Chain):
    def __init__(self):
        super(Model, self).__init__(
            l1=L.Linear(2, 1),
        )

    def __call__(self, x):
        h = F.sigmoid(self.l1(x))
        return h


model = Model()
optimizer = optimizers.MomentumSGD(lr=0.01, momentum=0.9)
optimizer.setup(model)

gpu_device = 0
cuda.get_device(gpu_device).use()
model.to_gpu(gpu_device)
xp = cuda.cupy

x = Variable(xp.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=xp.float32))
t = Variable(xp.array([[0], [1], [1], [1]], dtype=xp.float32))

for i in range(0, 3000):
    optimizer.zero_grads()
    y = model(x)
    loss = F.mean_squared_error(y, t)
    loss.backward()
    optimizer.update()

    print("loss: ", loss.data)

print(y.data)
```

ようは、np を cuda.cupy に置き換えるだけ。
あとは、以下の少しのおまじないを書くだけです。

```python
gpu_device = 0
cuda.get_device(gpu_device).use()
model.to_gpu(gpu_device)
xp = cuda.cupy
```

簡単！なのか？
==========

簡単！でこの記事を終わりとしたいのですが、これが簡単なのは、GPU側に送った行列をこちらに戻す必要が無いためです。つまり、to_gpu はありますが、to_cpu はここでは使って居ません。

といっても、教師あり学習の機械学習であれば、学習時は、to_gpu だけで済む事が多いと思います。to_cpu が必要なのは、学習した結果をとってきて、次の学習に生かす機械学習、、強化学習の Experimental Pool とかAuto encoderとかでしょうか。（このあたり勉強中であまり自信ありません。。）

強化学習をGPU化
==============

といわけで、強化学習をgpu化してみます。
なるべくシンプルなものが良いので、こちらを使わせてもらいました。 

https://github.com/dsanno/deep-learning-study/blob/master/src/open_ai_gym/train.py

とりあえず、np->xp
--------------

まず、単純に np -> xp してみます。

https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/646740c8d5afb811e815c4259205734053aac2b9

こんなエラーが出ます。

```
Traceback (most recent call last):
  File "train.py", line 208, in <module>
    main()
  File "train.py", line 196, in main
    update(agent, target_agent, optimizer, ex_pool, batch_size)
  File "train.py", line 122, in update
    indices = xp.random.permutation(available_size)[:batch_size]
AttributeError: 'module' object has no attribute 'permutation'
```

cupy は、[permutation](https://docs.scipy.org/doc/numpy/reference/generated/numpy.random.permutation.html) を持ってないようです。
cupyのドキュメントを見ても、[確かにありません](http://docs.chainer.org/en/stable/cupy-reference/random.html)。
これは、GPU側でやらなくてもいい処理なので、素直に、npで書くことにします。

```
    indices = np.random.permutation(available_size)[:batch_size]
```

dtypeが必要？
-----------

次に以下のようなエラーが出ます。（タイミング次第で別のエラーが出ることもある様です）

```
Traceback (most recent call last):
  File "train.py", line 208, in <module>
    main()
  File "train.py", line 196, in main
    update(agent, target_agent, optimizer, ex_pool, batch_size)
  File "train.py", line 125, in update
    state = xp.asarray(state)
  File "/usr/local/lib/python2.7/site-packages/cupy/creation/from_data.py", line 47, in asarray
    return cupy.array(a, dtype=dtype, copy=False)
  File "/usr/local/lib/python2.7/site-packages/cupy/creation/from_data.py", line 27, in array
    return core.array(obj, dtype, copy, ndmin)
  File "cupy/core/core.pyx", line 1538, in cupy.core.core.array (cupy/core/core.cpp:53535)
  File "cupy/core/core.pyx", line 1554, in cupy.core.core.array (cupy/core/core.cpp:53160)
ValueError: Unsupported dtype object
```

よくわかりませんが、エラーが出た箇所に、dtypeを以下のように、指定しておきます。

```python
    state = xp.asarray(state, dtype=xp.float32)
    action = xp.asarray(action, dtype=xp.int32)
    reward = xp.asarray(reward, dtype=xp.float32)
    next_state = xp.asarray(next_state, dtype=xp.float32)
    has_next = xp.asarray(has_next, dtype=xp.float32)
```

Advanced indexing
----------------

次に以下のようなエラーが出ます。

```
Traceback (most recent call last):
  File "train.py", line 208, in <module>
    main()
  File "train.py", line 196, in main
    update(agent, target_agent, optimizer, ex_pool, batch_size)
  File "train.py", line 133, in update
    y = reward + agent.gamma * has_next * target_agent(next_state).data[(six.moves.range(len(next_action))), next_action]
  File "cupy/core/core.pyx", line 1027, in cupy.core.core.ndarray.__getitem__ (cupy/core/core.cpp:22884)
ValueError: Advanced indexing is not supported
```

[Advaned indexing](https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing)が未サポートのようです。

これは、以下のようなコードの部分です。

```python
    y = reward + agent.gamma * has_next * target_agent(next_state).data[(six.moves.range(len(next_action))), next_action]
```

`target_agent(next_state).data[(six.moves.range(len(next_action))), next_action]` が問題なのです。Advanced indexingの指定で分かりにくいので、forを使った書き方に変えてみると、以下のような感じです。

```python
[a[i][next_action[i]] for i in six.moves.range(len(next_action))]
```

これを CuPy にやらせるのは難しいため、CUDAのコードを直接書きます。
以下が参考になります。

chainer.cuda.elementwiseを使ってGPUで処理を行う
http://qiita.com/dsanno/items/a68d5aaae0a542c0f388

同じ処理は以下のように書けます。

```python
    target_data = cuda.elementwise(
        'raw T x, S t',
        'T y',
        'int ind[] = {i, t}; y = x[ind];',
        'action_select_fwd',
    )(target_data, next_action)
```

ここまでの修正で、一通り動くようになりました。

https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/c4bf41abd590e0c6ee2b5d74c1296aefcbfceef5

一つ目からの差分は以下で見れます。

https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/revisions#diff-e44f4a60e89e820dde9bb27afa634965

遅い！！
=====

さて、これで動かしてみると、かなり遅くなりました。
これは、[このあたり](https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/c4bf41abd590e0c6ee2b5d74c1296aefcbfceef5#file-train-py-L122)で、CPU上のArrayからGPUへのコピーが発生しているためです。
xp.takeを使うように変えます。ExperimentalPoolにメソッドを追加します。

```python
    def take(self, indices):
        if self.pos < self.size:
            offset = 0
        else:
            offset = self.pos % self.size - self.size
        indices += offset
        return (xp.take(self.states, indices, axis=0), xp.take(self.actions, indices, axis=0),
                xp.take(self.rewards, indices, axis=0), xp.take(self.states, indices + 1, axis=0),
                xp.take(self.nexts, indices, axis=0))
```

全体はこちらです。
https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/86d6dabae832a9023d5274528028f33b5e09ddfc

差分はこちら
https://gist.github.com/ikeyasu/8c973ebde7cbd65301052d2226d8060a/revisions#diff-e44f4a60e89e820dde9bb27afa634965

まだ遅いよ！
=========

これで少し早くなったようですが、CPUと比較して、特に早くはありません。

* もともと計算量が多くないサンプルなので、GPUのボトルネックの方が多いのかも
* データ量は少なくても、CPUとGPUを行き来するところが残っているのでそこが悪いのか？
* GPUの性能が悪いのか？

等が考えられますが、そもそも計算量が少ないので、計算量が大きいサンプルで次は試してみたいと思います。（実は、既にCNNを使った強化学習のサンプルを作って試してみたのですが、まだ早くなっておらず、Advent Calendarに間に合いませんでした。。）

次は
===

@pankona さんです。（明日はもう、クリスマスイブですね）