Chainerとcupyをcuda版numpyとして使ってオートマトン

Numpyを使ったプログラムをcupy化しようとしたら、
計算量多そうなscipyのcorrelate2dがあってどうなるかと思いましたが、chainerの畳み込み層を使ってみたらなんとかなった話です。

下の方にプログラム全体をまとめてあり、これを適当なpythonスクリプトとして保存して実行できるかと思います。

# 環境

```
Ubuntu 16.04
Python 3.5.2
numpy 1.14.3
chainer 4.1.0
CUDA toolkit 9.1
OpenCV 3.4.3
```

# 準備

六角形のグリッドでrock-paper-scissorsオートマトンを実装したこちら：
https://qiita.com/sage-git/items/183999c1006999959d68
これを普通の方眼グリッドでやってみます。

## CPU版コード

まず普通のnumpyで組んだプログラムです。

```python
import random
import numpy as np
from scipy import signal
import cv2

MAX_HP = 18

cell_color = [(255, 255, 255), (0, 0, 0), (36, 120, 1), (1, 84, 240)]

neighbor = np.array([
    [1, 1, 1],
    [1, 0, 1],
    [1, 1, 1]
], dtype=int)

class RockPaperScissorAuto(object):
    def __init__(self, W, H):
        self.W = W
        self.H = H
        self.F = np.zeros((self.H, self.W), dtype=int)
        self.HP = np.ones_like(self.F)
    
    def step(self):
        empty = self.F == 0
        rock = self.F == 1
        paper = self.F == 2
        scissor = self.F == 3
        nxp = signal.correlate2d(paper, neighbor, mode="same")*(rock + empty)
        nxs = signal.correlate2d(scissor, neighbor, mode="same")*(paper + empty)
        nxr = signal.correlate2d(rock, neighbor, mode="same")*(scissor + empty)
        self.HP -= (nxp + nxs + nxr)
        mut = self.HP < 1
        self.F[mut] += 1  
        self.F[self.F == 4] = 1
        self.HP[mut] = MAX_HP 
    
    def mutation(self, rate=0.0005):
        mut = (np.random.rand(self.H, self.W) < rate)
        self.F[mut == 1] += 1
        self.F[self.F == 4] = 1
        self.HP[mut == 1] = MAX_HP

    def make_image(self, cell_rad=10):
        width = self.W
        height = self.H
        img = np.zeros((height, width, 3), dtype=np.uint8)
        for i in range(4):
            img[self.F == i] = cell_color[i]
        return img

def main():
    F = RockPaperScissorAuto(480, 320)
    waiting = 10
    rate = 0.01
    while True:
        F.step()
        F.mutation(rate=rate)
        rate *= 0.99
        img = F.make_image(cell_rad=5)
        cv2.imshow("test", img)
        r = cv2.waitKey(waiting)
        if r == ord("+"):
            waiting = max(10, waiting//2)
        if r == ord("-"):
            waiting = min(1000, waiting*2)
        if r == ord("q"):
            break
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()
```

## 六角形グリッドのときからの変更点

### mutation

六角形グリッドの時と比べて、`mask`がなくなりました。

### step

あるセルの周りのセルをカウントしてから、それに基づいて変異させます。

ふと、「`HP`が１以下になったものが変異する」とすると、`mutation`と同じルーチンで世代交代を表現できることに気づきました。なので後半3行は`mutation`とほとんど変わりません。この変更は前回の六角形グリッドのものでも有効です。

### make_image

今回は丸を並べる必要がないので、セルをそのままピクセルに割り当てます。
`img`の`shape`が`(H, W, 3)`なのに対し、`self.F == i`の`shape`は`(H, W)`です。これでも`img[self.F == i]`というようにindexに入れることができて、余った`shape`の次元である3に対して`cell_color[i]`を代入しています。さすがnumpyのブロードキャスト。

### cell_color

軍産複合体（参考：https://qiita.com/sage-git/items/183999c1006999959d68 のネタばれ）が導入した赤、青、黄の３色は目が痛くなったので、落ち着いた色に変えます。いろいろ探した結果、歌舞伎の定式幕になりました。

### 変異確率

`mutation`に渡す`rate`をステップごとに減衰させることで、初期に激しく揺動して渦を作り、後半は静かな環境で安定した渦を観察します。

これで実行してみた結果です。

![test.gif](https://qiita-image-store.s3.amazonaws.com/0/258019/1d017cdc-e1a8-d4ae-f7cb-f88eec0e1c29.gif)

典型的なBZ反応のシミュレーションができたかと思います。

# cupy化

## import

```python
import numpy as np
import cupy as xp
```
openCVで画像表示をするので、numpyも使います。


## initialize

`np.zeros()`など、`numpy`で作る配列をそのまま`xp.zeros()`などのように`cupy`に変えていくだけです。

## correlateをconvolution2dで

[Python：Scipyのcorrelate2D (convolve2d)でいろいろ数える](https://qiita.com/sage-git/items/1b31c88bfaa9acb77daa)
こちらで触れたように、correlateとconvolutionは場合によっては同じように使えます。
Chainerのconvolutionといえば畳み込み層。ということで`convolution_2d`を使ってみます。
え？ニューラルネットのConv層は数学的には厳密には畳み込みではない？大丈夫です、数学的なことは問題になりませんでした。それにうまく行かなかったらこの記事を書いていません。

correlate2dで隣接マスの定義として使っていた行列をそのままConv層の$W$とします。
ただし、`float`でないといけないので、とりあえず半精度浮動小数点にしました。
また、$W$の次元は$(c_I, c_O, h_K, w_K)$で、$c_I$と$c_O$は入出力チャンネル数です。今回はどちらも１なので、単純に`reshape`で`1`を増やしておきます。

```python
neighbor = xp.array([
    [1, 1, 1],
    [1, 0, 1],
    [1, 1, 1]
], dtype=xp.float16).reshape(1, 1, 3, 3)
```

まずrockなどセルの状態の数え方から。現在rockのセルが`1`、それ以外が`0`の行列を作ります。

```python
    rock = xp.array(self.F == 1, dtype=xp.float16).reshape(1, 1, self.H, self.W)
    paper = xp.array(self.F == 1, dtype=xp.float16).reshape(1, 1, self.H, self.W)

```
これも`float`にします。また、convolution2dの入力は$(N, c, w, h)$ですが、バッチ数$N$とチャンネル数$c$は１として`reshape`します。

続いて、correlate2dと同値の関数を使って、あるセルの周りのpaperの数を数えます。

```python
import chainer.functions as F
v = F.convolution_2d(paper, neighbor, pad=1)
```

なお、`convolution_2d`は`chainer.Variable`を返します。これは余計なデータを持っていますので、とっとと`int`型のcupy配列になおしておきます。
ついでに、周りのpaperの数を気にするのはrockかemptyだけなので、それでマスクします。

```python
    nxp = xp.array((v*(rock + empty)).array, dtype=int).reshape(self.H, self.W)
```
この`nxp`は、rockまたはemptyの周りのpaperの数で、この数がそのまま`HP`を減らす量になります。
ここでrock, paper, scissorsに同様の処理をするのですが、この`convolution2d`は並列にできそうです。ということはミニバッチ化できそうですね。

## ミニバッチ化

もう一度数えるところから。

```python
    rock = (self.F == 1).reshape(1, self.H, self.W)
    paper = (self.F == 2).reshape(1, self.H, self.W)
    scissor = (self.F == 3).reshape(1, self.H, self.W)
```

convolution2dの入力は$(N, c, w, h)$ですが、バッチ数$N$ 用の次元が`reshape`から減りました。これを

```python
    rps = xp.array(xp.stack((rock, paper, scissor)), dtype=xp.float16)
```

こうしてミニバッチ化します。`stack`によりミニバッチ用の次元が追加されます。ミニバッチ中のサンプル数は３です。これを`convolution_2d`に渡し、cupyの配列に変換しつつミニバッチを分解します。

```python
    cr, cp, cs = xp.vsplit(xp.array(F.convolution_2d(rps, neighbor, pad=1).array, dtype=int), [1, 2])
```

あとは必要な領域だけにmaskします。

```python
    nxp = (cp*(rock + empty)).reshape(self.H, self.W)
    nxs = (cs*(paper + empty)).reshape(self.H, self.W)
    nxr = (cr*(scissor + empty)).reshape(self.H, self.W)
```

ここからの処理（HP減とセルの種類変化）はNumpy版と同じです。

## image化

cupyの配列はそのままではcv2に渡せないので、GPUからCPUに渡してやります。

```python
    img = xp.zeros((height, width, 3), dtype=xp.uint8)
    # 中略、imgに対して必要な描画を行う
    return np.array(cuda.to_cpu(img))
```

# 実行

## プログラム全体

以上をまとめると、以下のようなプログラムになります。

```python:cupy_grid_automaton.py
import random
import numpy as np
import cupy as xp
from chainer import cuda
import chainer.functions as F
import chainer.variable as V
import cv2

MAX_HP = 12

cell_color = [(255, 255, 255), (0, 0, 0), (36, 120, 1), (1, 84, 240)]

neighbor = xp.array([
    [1, 1, 1],
    [1, 0, 1],
    [1, 1, 1]
], dtype=xp.float16).reshape(1, 1, 3, 3)

class RockPaperScissorAuto(object):
    def __init__(self, W, H):
        self.W = W
        self.H = H
        self.F = xp.zeros((self.H, self.W), dtype=int)
        self.HP = xp.ones_like(self.F)
    
    def step(self):
        empty = (self.F == 0).reshape(1, self.H, self.W)
        rock = (self.F == 1).reshape(1, self.H, self.W)
        paper = (self.F == 2).reshape(1, self.H, self.W)
        scissor = (self.F == 3).reshape(1, self.H, self.W)
        rps = xp.array(xp.stack((rock, paper, scissor)), dtype=xp.float16)
        cr, cp, cs = xp.vsplit(xp.array(F.convolution_2d(rps, neighbor, pad=1).array, dtype=int), [1, 2])
        nxp = (cp*(rock + empty)).reshape(self.H, self.W)
        nxs = (cs*(paper + empty)).reshape(self.H, self.W)
        nxr = (cr*(scissor + empty)).reshape(self.H, self.W)

        self.HP -= (nxp + nxs + nxr)
        mut = self.HP < 1
        self.F[mut] += 1  
        self.F[self.F == 4] = 1
        self.HP[mut] = MAX_HP 
        
    def mutation(self, rate=0.0001):
        mut = (xp.random.rand(self.H, self.W) < rate)
        self.F[mut == 1] += 1
        self.F[self.F == 4] = 1
        self.HP[mut == 1] = MAX_HP

    def make_image(self, cell_rad=10):
        width = self.W
        height = self.H
        img = xp.zeros((height, width, 3), dtype=xp.uint8)
        for i in range(4):
            img[self.F == i] = cell_color[i]
        return np.array(cuda.to_cpu(img))
    

def main():
    F = RockPaperScissorAuto(640, 480)
    waiting = 10
    rate = 0.02
    while True:
        F.step()
        F.mutation(rate=rate)
        rate *= 0.99
        img = F.make_image(cell_rad=5)
        cv2.imshow("test", img)
        r = cv2.waitKey(waiting)
        if r == ord("+"):
            waiting = max(10, waiting//2)
        if r == ord("-"):
            waiting = min(1000, waiting*2)
        if r == ord("q"):
            break
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()
```

## パフォーマンス

系のサイズをいろいろ変えた時のGPUの使用率です。
GPUの使用率については、実行中に`nvidia-smi`コマンドを１０回程度打って、僕のインスピレーションによる平均を取ったものです。
GPUはGTX1060です。

![gpuusage.png](https://qiita-image-store.s3.amazonaws.com/0/258019/edebae81-9612-6b5c-c058-7fbf9f043a0c.png)

右端のは確か2560x1960です。

これでnumpyのGPU化が必要であればcupyで、そして適宜chainerを使えばscipyのgpu化もできるのかもしれません。

Tensorflowでも同様にできる気がします。

# その他

オートマトンの見た目ですが、一言、キモい。