PyTorch で拡散シミュレーション

なぜこんなことを・・・？

• 拡散のシミュレーションがやりたい！
• PyTorch を使えるようになりたい！

という２つの動機の単純な帰結としてそうなった・・・わけではなく、実はすごく理にかなった選択なのです。・・・本当です。信じてください。

ともあれ、終わってみて考えると、 PyTorch の基本に一通り触れることができ、チュートリアルとして良さそうに思います。私も始めたばかりなのですが、興味がある人はぜひ見てみてください。

• PyTorch でのデータの扱い方
• GPU の使い方
• 畳み込み
• 結果が見れてわかりやすい
• なんと動画までつくれる

アイデア

インクでイメージ

$t=0$ で水中にインクを落としたとします。次の時刻 $t=1$ では、インクは拡散し、元の位置よりも広がりをもって分布します。

では、もし２ヶ所にインクを落としたらどうなるでしょうか？　時間が経つと、それぞれが同じように広がります。

水中全体にインクがあっても全く同じように捉えることができます。時間が経つと、それぞれの位置でインクが周りに広がり、結果としてインクの濃度分布が変化します。

結局どうやるの？

拡散のシミュレーションをするにあたって考えねばならないことが２つあります。

まず、インクがどのように広がるかということですが、実は正規分布と同じ形になります。

拡散方程式の解と粒子数の時間変化
http://statmodeling.hatenablog.com/entry/diffusion-equation-1

ある場所でのインク分子は、次の時刻ではベルカーブ状に周囲に分配されます。逆に、ある場所でのインク分子の量は、前の時刻での周囲のインク分子をベルカーブの重みでたし合わせたものになります。

次に、すべての点で濃度変化を計算しなければなりませんが、計算方法を単純に表現すると、「インク濃度に重みをかけて足し合わせる」ことになります。これは畳み込み演算そのもので、 PyTorch などの深層学習ライブラリの得意とするものです。

実装時の Tips

コードはここにあるので、気になる人はどうぞ
https://nbviewer.jupyter.org/github/Lirimy/simulation/blob/master/diffusion.ipynb (notebook viewer)
修正していて、以下の記述と異なるところもあります

CPU と GPU の切り替え方

コメントアウトでデータ型を設定しておいて、テンソル作成時に適用すればよい

#dtype = torch.FloatTensor # CPU
dtype = torch.cuda.FloatTensor # GPU

tensor = torch.Tensor(...).type(dtype)

NumPy データで畳み込みするには？

PyTorch での２次元畳み込みで利用するデータ形式は4次元 Variable
そこで、NumPy データを順番に変換していく

1. NumPy 行列 x （２次元）をつくる
2. [[x]]で4次元にする
3. torch.Tensor(...) でTensor型にする
4. Variable(...) でラップする

import numpy as np
import torch
from torch.autograd import Variable

x = np.zeros((10, 10)) # 2-D numpy matrix
x = Variable(torch.Tensor([[x]]).type(dtype), requires_grad=False)

ちなみに増えた２次元は、

• バッチ内でのデータ番号
• チャンネル

今回はどちらも使わない

GPU での計算結果にアクセスするには？

上でやったことの逆をたどる

1. _.data で Variable 内の Tensor にアクセス
2. _.cpu() で GPU からデータを転送
3. _.numpy() で Tensor から NumPy に変換
4. _[0, 0] で４次元（ダミー）から２次元に

この操作で2次元 NumPy 行列が得られる

x # Variable containing torch.cuda.FloatTensor
x_np = x.data.cpu().numpy()[0, 0]

パディングオプションでの注意

畳み込み F.conv2d のオプション padding=m としたときに、 m は厳密に int でなければなりません。
float などもってのほか、 np.int ですら許してもらえません。

import torch.nn.functional as F

m = int(m) # DO THIS
y = F.conv2d(x, w, padding=m)

シミュレーションの結果

拡散の様子（アニメーション）

畳み込みで０パディングしているため、境界から媒質が流出している

参考

matplotlibでアニメーション
https://qiita.com/chez_sugi/items/93ff2efad50f16bc8c37

matplotlibのanimation.FuncAnimationを用いて柔軟にアニメーション作成
https://qiita.com/AnchorBlues/items/3acd37331b12e844d259

Jupyterでmatplotlibのjavascriptアニメーションを動かす
https://qiita.com/yoku_001/items/11ac4ad1df349095655c

パフォーマンス測定

CPU と GPU のパフォーマンスの違いを知りたい

2次元畳み込みに要する実時間を測定した
Jupyter の %%timeit がとても便利

条件

• グリッドサイズ $n=200$
• 変換行列サイズ $2m+1=47$
• パディング $m=23$
• 繰り返し 100 回

$(246, 246)$ 行列と $(47, 47)$ 行列を畳み込むことになる
$ m=23 $ が中途半端なのは、標準偏差 $ \sigma=10 $ から自動的に決まった値だから

環境

• i5-2500K
• GTX 980

結果

• CPU 9-10 sec
• GPU 6-10 msec

GPU のほうが 1000 倍速い！

まとめ

• 正規分布の畳み込みで拡散シミュレーションができる！
• 畳み込みといったら PyTorch !!! （別に TensorFlow とかでもよい）
• GPU はすごい！
• アニメーションを見るのはたのしー！

編集履歴

2018.1.30 初版
2018.2.2 Javascript アニメーションを追加