Chainer チュートリアルも公開されたし GPU がもうちょっと楽しくなる CuPy カーネル入門

最近教育にも力を入れてるなーと思っていたら Chainer チュートリアルも公開されて。ディープラーニング入門なところを日本語でちゃんと勉強できるコンテンツは今すごく求められていますね。

(僕も一通り見ましたが、自分の Chainer のコードの参考にさせて頂きました)

tutorials.chainer.org

ということで、今回はこの第10章 CuPy 入門 + α でも書こうかと。

10. CuPy 入門 — ディープラーニング入門：Chainer チュートリアル

CuPy は PFN さんが開発している、NumPy っぽく使えるのに GPU で単純な計算を超高速化するライブラリ。実は Chainer とは独立でも使えるので、非常に汎用性が高い。(著作権があるので雰囲気だけ伝えると、こんな感じのロゴを見たらそれである)

cupy.chainer.org

普通はチュートリアルにあるように cupy を numpy っぽく使えばいいんだけど、自分で作りたいなーと思ったときに登場するのが CuPy の CUDA Kernel クラス。色んなコードで見るけど、とっつきにくくてよく分からないなぁ… と思っていたので調べてみるとめちゃめちゃ簡単だったというのをまとめようと思います。

直接書いちゃう Raw は置いとくとして、あとは Elementwise と Reduction の2つ。

https://docs-cupy.chainer.org/en/stable/tutorial/kernel.html

Elementwise

まず Elementwise は、例えばベクトル+ベクトルの要素の足し算を並列化するような計算。例えば、ある位置 x から、速度 v で時間 dt 移動して x + v*dt という位置に移動する計算をしてみる。簡単に N 次元に拡張できるので、全部の次元は独立に計算ができるとする。この時の CuPy カーネルは、

import cupy as cp
ew_kinetic_kernel = cp.ElementwiseKernel(
    'T x, T v, T dt', # 入力
    'T y',            # 出力
    'y = x + v*dt',   # 計算式
    'ew_kinetic')     # 名前

引数は、入力、出力、計算式、名前。型指定してもいいけど “T” で OK ジェネリック。で、呼ぶときは x, v, dt を引数にして、y で受ければ要素ごとに x+v*dt をしてくれる。こんな感じ。

x = [-10, 0, 10]
v = [100, 200, 300]
x_cp = cp.asarray(x).astype(cp.float32)
v_cp = cp.asarray(v).astype(cp.float32)

y = kinetic_kernel(x_cp, v_cp, 1)

Chainer チュートリアルっぽく N を変化させて Numpy、CuPy そのまま、Kernel の実行時間を比較してみました (適当に x[n] = cos(pi(1-2*n/N)), v[n] = sin(pi(1-2*n/N)) ってな感じで)。

N が少ないとむしろ遅くなるんですが、N が多くなると圧倒的な計算パフォーマンス (しかも cupy そのまま使うより早い)。

Reduction

次は Reduction。これは例えば要素毎の和をとるような計算。例えば、さっきみたいな簡単な力学から、質量 m の物体が速度 v で動いている運動エネルギーを計算する。すると運動エネルギー K = 1/2 m v^2 である。N次元だとこの和を計算すればいい。このときのカーネルは、

rd_kinetic_kernel = cp.ReductionKernel(
'T m, T v',        # 入力
'T y',             # 出力
'v * v',           # 前処理
'a + b',           # 繰り返し計算
'y = 0.5 * m * a', # 計算式(後処理)
'0',               # 初期値
'rd_kinetic')      # 名前

これがまた分かりにくい…。入力、出力、計算式、名前まではさっきと同じである。

じゃあ中ではどういう計算をするか？順序とかいいから、とにかく K を足せばいい、つまりそういう感じ。まず v を要素毎に二乗する (前処理; Elementwise)。そして、a = 0 (初期値)、からのひたすら a = a + b (繰り返し計算; ここで b が前処理した要素である)。つまり K_1 + K_2 + …… K_N をひたすらしまくる。そして最後に 0.5 * m * a をして y に渡すので、求めていた計算結果が得られる。呼び出し方はさっきと同じである。計算結果が以下。