　コードの中身はどうでもいいからGPUでどれくらい速くなるか知りたいという方は下までスキップして下さい．

環境

　ツクモで17万くらいで購入したゲーミングPCを使います．

OS: Windows 10 Home
CPU: Intel Core i9-9900KF
RAM: 16.0 GB
GPU: GeForce RTX 2070 SUPER
VRAM: 8.0 GB
Python 3.8.2
numpy 1.19.1
scipy 1.5.2
PyTorch 1.9.0+cu111
CUDA 11.1

目的

　scipy.sparseだと計算時間がかかりすぎて効率が悪かったのでGPUが使えるPyTorchで高速化したかった．

GPU環境の（再）構築

　久々に使おうとしたらつまづいたので簡単にメモ．初めて導入するのであればVisual Stadioが必要なのでこちらも参考．

torch.cuda.is_available()がTrueなのを確認．
Falseなら環境の（再）構築が必要．以前に構築済みでダメになっているのならドライバの更新時にCUDAとバージョンが合わなくなったと思われる．
nvcc -VでCUDAのバージョンを確認する．Pytorchの公式インストールガイドのバージョンと一致するか確かめる．
不一致の場合，念のために古いCUDAとcuDNNを削除しておく．コントロールパネルから可能．
必要なバージョンのCUDAをNVIDIA公式から取ってくる．安易に取ってくると最新版となりPyTorchと合わなかったのでアーカイブから選ぶこと．
CUDAに合わせてcuDNNをダウンロード．こちらもアーカイブから選ぶ．
cuDNNをパスが通ってる場所に置く．
PyTorchインストールガイドのコマンドをコピペして実行．最初の'torch.cuda.is_available()`が'True'になることを確認．

問題設定

　以前投稿したように2次元格子点をバネで繋いだ時の運動をルンゲクッタ法で解きます．密行列で計算すると一辺の原子数$N$に対して$O \left( N^4 \right)$の要素数になるはずです．疎行列で$O \left( N^2 \right)$のはず．分子の速度（というか平均運動エネルギー）で温度を議論するにはボルツマン分布を再現できている必要があるので原子数はなるべく増やしたいですが，この通りちょっと原子数を増やすと計算量があっといまに膨れ上がってしまいます．こうなると時間はかかるしメモリに乗り切らないしで困りました．（原子数個の速度で温度を議論しているような文献には身構えて下さい．分布の幅は結構広いです．「その速度は違う温度でも〇%の確率で得られますが，その議論は本当に正しいのですか？」と言えるとかっこいいですね．）
　さて計算速度を向上させる方法は主に2つあります．1つ目がより高速なマシンを使うこと，2つ目が計算量を減らすことですね．メモリ問題にしたって同じです．アプローチをハードとソフトのどちらからかけるかです．今回はハード面はGPUの活用，ソフト面は疎行列の活用です．・・・・・・と言いながらソフト面は以前の投稿でも既に実行済みなので，改めて密行列の場合も計算してみて比較します．
　なお残念ながらGPUを活用した場合に疎行列がうまく扱えなかったため，そちらを探されている方は諦めて下さい．これはPyTorchにしてもTensorFlowにしても同じようで，スパーステンソルはbroadcastが定義されていなかったり，通常のテンソルとのアダマール積が定義されていなかったためです．従ってタイトルの通りゲーミングPCを使う場合は簡単にVRAMが足りなくなります．行列積は使わないのでブロックごとにGPUに投げればいいだけかもしれませんが試してません．

解法

　まずCPUで密行列を計算するコードです．長くなるのでコメントは全消しです．説明が必要であれば以前の投稿を確認願います．

cpu_dense.py

from pathlib import Path
import numpy as np
from matplotlib import pyplot as plt

class RungeKutta(object):
    def __init__(self, dt, grad):
        self.dt = dt
        self.grad = grad

    def calc_k(self, t, y, **kwargs):
        k1 = self.grad(t, y, **kwargs)
        k2 = self.grad(t + self.dt/2, y + self.dt/2*k1, **kwargs)
        k3 = self.grad(t + self.dt/2, y + self.dt/2*k2, **kwargs)
        k4 = self.grad(t + self.dt, y + self.dt*k3, **kwargs)
        return k1, k2, k3, k4

    def calc(self, t, y, **kwargs):
        k1, k2, k3, k4 = self.calc_k(t, y, **kwargs)
        y_new =  y + self.dt/6*(k1 + 2*k2 + 2*k3 + k4)
        t_new = t + self.dt
        return t_new, y_new


def hooke(t, v, **kwargs):
    X = kwargs['X']
    n = X.shape[0]
    R = []
    r = 0
    for i in range(n):
        xi = kwargs['C']*(X[i].reshape(-1,1))
        xj = kwargs['C']*(X[i].reshape(1,-1))
        R.append(xi-xj)
        #r += (xi-xj).power(2)
        r += (xi-xj)**2
    L = kwargs['C']*(kwargs['L'])
    r = np.sqrt(r)
    f = kwargs['K']*(r-L)

    F = []
    for i in range(n):
        F.append(
            np.nan_to_num(f*(R[i])/r)
                .sum(axis=-1)
                .reshape(X[i].shape)
                /kwargs['M']
            )
    F = -np.array(F)

    return F

def velosity(t, x, **kwargs):
    return kwargs['V']
 
def boundary(t, phi, a, k, m):
    omega = np.sqrt(k/m)
    x = np.cos(omega*t + phi)*a*np.array([[0.5, 1]])
    vx = -omega*np.sin(omega*t + phi)*a*np.array([[0.5, 1]])
    return x, vx

def main(Nx=5, Ny=3, dt=1.e-16, upto=1.e-12,
        k=1.e3, l=1.e-10, m=1.e-25, a=3.e-12):
    t = 0
    I = np.ones(Nx*Ny).reshape(Ny,Nx)
    E = (np.eye(Nx*Ny))

    Iu = I.copy()
    Iu[0,:]=0
    Iu = Iu.reshape(-1,1)

    Il = I.copy()
    Il[:,0] = 0
    Il = Il.reshape(-1,1)

    Ir = I.copy()
    Ir[:,-1] = 0
    Ir = Ir.reshape(-1,1)

    Id = I.copy()
    Id[-1,:] = 0
    Id = Id.reshape(-1,1)

    C = (np.roll(E*Iu, -Nx, axis=1)
            + np.roll(E*Il, -1, axis=1)
            + np.roll(E*Ir, 1, axis=1)
            + np.roll(E*Id, Nx, axis=1)
            )
    K = C*k

    L = C*l

    M = np.ones(Nx*Ny).reshape(Ny,Nx)*m

    x = np.arange(1,Nx+1)*l
    y = np.arange(Ny,0,-1)*l
    x, y = np.meshgrid(x,y)
    xb0 = x.copy()
    yb0 = y.copy()
    dx = np.random.randn(Ny,Nx)*l/300
    dy = np.random.randn(Ny,Nx)*l/300
    dx[:,0] = 0
    dx[:,-1] = 0
    x = x + dx
    y = y + dy
    phi = np.random.rand(Ny*2).reshape(Ny, 2)*2*np.pi
    xb, vb = boundary(0, phi, a, k, m)
    x[:,0] += xb[:,0]
    x[:, -1] += xb[:,1]
    x = np.stack([x, y])
    v = np.zeros_like(x)
    v[0,:,0] += vb[:,0]
    v[0,:,-1] += vb[:,1]


    V = RungeKutta(dt, hooke)
    X = RungeKutta(dt, velosity)

    folder = Path(
        f'dense_k{k}_l{l}_m{m}_a{a}_dt{dt}'
        )
    folder.mkdir(exist_ok=True)
    wd = folder/f'Nx{Nx}_Ny{Ny}_time{upto}'
    wd.mkdir(exist_ok=True)
    res_t = open(wd/f'time.txt', 'w')
    res_vx = open(wd/f'velosity_x.txt', 'w')
    res_vy = open(wd/f'velosity_y.txt', 'w')
    res_x = open(wd/f'x.txt', 'w')
    res_y = open(wd/f'y.txt', 'w')

    fig, ax = plt.subplots()
    ax.scatter(x[0], x[1], c='k')
    fig.savefig(wd/'Init.png', bbox_inches='tight')

    np.savetxt(res_t, np.array([t]))
    np.savetxt(res_vx, v[0])
    np.savetxt(res_vy, v[1])
    np.savetxt(res_x, x[0])
    np.savetxt(res_y, x[1])

    for txt in [res_t, res_vx, res_vy, res_x, res_y]:
        txt.write('\n')
    while t<=upto:
        t_new, v_new = V.calc(
                    t, v, X=x, K=K,
                    L=L, M=M, C=C
                )
        xb ,vb = boundary(t_new, phi, a, k, m)
        xb[:,0] += xb0[:,0]
        xb[:,-1] += xb0[:, -1]
        v_new[0, :, 0] = vb[:,0]
        v_new[0, :, -1] = vb[:,1]
        v_new[1, :, 0] = 0
        v_new[1, :, -1] = 0

        t_new, x_new = X.calc(t, x, V=v)
        x_new[0, :, 0] = xb[:,0]
        x_new[0, :, -1] = xb[:,1]
        x_new[1, :, 0] = yb0[:, 0]
        x_new[1, :, -1] = yb0[:, -1]

        t = t_new
        v = v_new
        x = x_new
        
        np.savetxt(res_t, np.array([t]))
        np.savetxt(res_vx, v[0])
        np.savetxt(res_vy, v[1])
        np.savetxt(res_x, x[0])
        np.savetxt(res_y, x[1])

        for txt in [res_t, res_vx, res_vy, res_x, res_y]:
            txt.write('\n')

    res_t.close()
    res_vx.close()
    res_vy.close()
    res_x.close()
    res_y.close()

    return wd

　続いてCPU上で疎行列を使うコードです．変更があるhooke関数は下記の通りです．main関数はnp.rollを使ってCを定義した後にC = csr_matrix(C)を追加してあります．というか以前投稿した記事のコードはこっちです．

cpu_sparse.py

from scipy.sparse import csr_matrix

def hooke(t, v, **kwargs):
    X = kwargs['X']
    n = X.shape[0]
    R = []#scipy.sparse CANNOT handle tensor.
    r = 0
    for i in range(n):
        xi = kwargs['C'].multiply(X[i].reshape(-1,1))
        xj = kwargs['C'].multiply(X[i].reshape(1,-1))
        R.append(xi-xj)
        r += (xi-xj).power(2)
    r = r.sqrt()
    f = kwargs['K'].multiply(r-kwargs['L'])

    F = []
    for i in range(n):
        F.append(
            np.nan_to_num(f.multiply(R[i])/r)
                .sum(axis=-1)
                .reshape(X[i].shape)
                /kwargs['M']
            )
    F = -np.array(F)

    return F

　最後にPyTorchでGPUを使うコードです．なんとかしてスパーステンソルを使おうとしましたが諦めました．普通のテンソルであればbroadcastは問題なく機能してくれます．
　class RungeKuttaはやはり変更なしなので省略です．ほとんどnumpy-likeに書けるので楽ですが，随所にGPUを使うための記述と倍精度を指定するための記述があります．torch.tensorの実数は単精度が標準になっていますが，numpyから変換した場合は倍精度になります．最大の注意点はnp.meshgridはindexing=xyが基本になってますが，torch.meshgridではindexing=ijでしか出力できないことです（パラメータが存在しません）．

gpu_dense.py

import torch

def hooke(t, v, **kwargs):
    X = kwargs['X']
    n = X.shape[0]
    xi = kwargs['C'].multiply(
            X.reshape(2,kwargs['C'].shape[0],1)
            )
    xj = kwargs['C'].multiply(
            X.reshape(2,1,kwargs['C'].shape[0])
            )

    R = xi-xj
    del xi, xj
    r = R**2
    r = r.sqrt()
    
    f = kwargs['k']*kwargs['C']*(r-kwargs['l'])

    F = -(torch.nan_to_num(f*R/r)
            .sum(axis=-1)
            .reshape(X.shape)
            /kwargs['m']
        )

    return F

def velosity(t, x, **kwargs):
    return kwargs['V']
    
def boundary(t, phi, a, k, m):
    device = torch.device(
                    'cuda:0' if torch.cuda.is_available()
                    else 'cpu'
                )
    omega = np.sqrt(k/m)
    x = (torch.cos(omega*t + phi)*
            a*torch.tensor([0.2, 1],device=device)
            )
    vx = (-omega*torch.sin(omega*t + phi)*
            a*torch.tensor([0.2, 1],device=device)
            )
    return x, vx

def main(Nx=5, Ny=3, dt=1.e-16, upto=1.e-12,
        k=1.e3, l=1.e-10, m=1.e-25, a=3.e-12):
    device = torch.device(
                    'cuda:0' if torch.cuda.is_available()
                    else 'cpu'
                )
    print('device:', device)

    t = 0
    E = torch.eye(Ny*Nx, device=device, dtype=torch.float64)

    Iu = torch.ones(
                Ny*Nx, device=device, dtype=torch.float64
            ).reshape(Ny,Nx)
    Iu[0,:]=0
    Iu = Iu.reshape(-1,1)

    Il = torch.ones(
                Ny*Nx, device=device, dtype=torch.float64
            ).reshape(Ny,Nx)
    Il[:,0] = 0
    Il = Il.reshape(-1,1)

    Ir = torch.ones(
                Ny*Nx, device=device, dtype=torch.float64
            ).reshape(Ny,Nx)
    Ir[:,-1] = 0
    Ir = Ir.reshape(-1,1)

    Id = torch.ones(
                Ny*Nx, device=device, dtype=torch.float64
            ).reshape(Ny,Nx)
    Id[-1,:] = 0
    Id = Id.reshape(-1,1)

    C = (torch.roll(E*Iu, -Nx, dims=1)
            + torch.roll(E*Il, -1, dims=1)
            + torch.roll(E*Ir, 1, dims=1)
            + torch.roll(E*Id, Nx, dims=1)
            )
    #C = torch.stack([C for i in range(len(X.shape[0]))]).to_sparse()

    x = np.arange(0,Nx)*l
    y = np.arange(Ny-1,-1,-1)*l
    #unlike np.meshgrid, torch.meshgrid afford tensor as indexing='ij' only
    x, y = np.meshgrid(x, y, indexing='xy')
    x = torch.from_numpy(x).to(device)
    y = torch.from_numpy(y).to(device)
    xb0 = torch.arange(0,Nx,device=device)*l
    yb0 = torch.arange(Ny-1,-1,-1,device=device)*l
    dx = torch.randn(Ny,Nx,device=device)*l/10
    dy = torch.randn(Ny,Nx,device=device)*l/10
    dx[:,0] = 0
    dx[:,-1] = 0
    x = x + dx
    y = y + dy
    phi = torch.rand(
            Ny*2,device=device
            ).reshape(Ny, 2)*2*np.pi
    xb, vb = boundary(0, phi, a, k, m)
    x[:,0] += xb[:,0]
    x[:, -1] += xb[:,1]
    x = torch.stack([x, y])
    v = torch.zeros_like(x)
    v[0,:,0] += vb[:,0]
    v[0,:,-1] += vb[:,1]

    V = RungeKutta(dt, hooke)
    X = RungeKutta(dt, velosity)

    folder = Path(
        f'pytorch_k{k}_l{l}_m{m}_a{a}_dt{dt}'
        )
    folder.mkdir(exist_ok=True)
    wd = folder/f'Nx{Nx}_Ny{Ny}_time{upto}'
    wd.mkdir(exist_ok=True)
    res_t = open(wd/f'time.txt', 'w')
    res_vx = open(wd/f'velosity_x.txt', 'w')
    res_vy = open(wd/f'velosity_y.txt', 'w')
    res_x = open(wd/f'x.txt', 'w')
    res_y = open(wd/f'y.txt', 'w')

    fig, ax = plt.subplots()
    ax.scatter(x[0].to('cpu'), x[1].to('cpu'), c='k')
    fig.savefig(wd/'Init.png', bbox_inches='tight')

    np.savetxt(res_t, np.array([t]))
    np.savetxt(res_vx, v[0].to('cpu').numpy())
    np.savetxt(res_vy, v[1].to('cpu').numpy())
    np.savetxt(res_x, x[0].to('cpu').numpy())
    np.savetxt(res_y, x[1].to('cpu').numpy())

    for txt in [res_t, res_vx, res_vy, res_x, res_y]:
        txt.write('\n')
    while t<=upto:
        t_new, v_new = V.calc(
                    t, v, X=x, k=m,
                    l=l, m=m, C=C
                )
        xb ,vb = boundary(t_new, phi, a, k, m)
        xb[:,0] += xb0[0]
        xb[:,-1] += xb0[-1]
        v_new[0, :, 0] = vb[:,0]
        v_new[0, :, -1] = vb[:,1]
        v_new[1, :, 0] = 0
        v_new[1, :, -1] = 0

        t_new, x_new = X.calc(t, x, V=v)
        x_new[0, :, 0] = xb[:,0]
        x_new[0, :, -1] = xb[:,1]
        x_new[1, :, 0] = yb0
        x_new[1, :, -1] = yb0

        t = t_new
        v = v_new
        x = x_new
        
        np.savetxt(res_t, np.array([t]))
        np.savetxt(res_vx, v[0].to('cpu').numpy())
        np.savetxt(res_vy, v[1].to('cpu').numpy())
        np.savetxt(res_x, x[0].to('cpu').numpy())
        np.savetxt(res_y, x[1].to('cpu').numpy())

        for txt in [res_t, res_vx, res_vy, res_x, res_y]:
            txt.write('\n')

    res_t.close()
    res_vx.close()
    res_vy.close()
    res_x.close()
    res_y.close()

    return wd

　あとはこいつらに縦横の原子数を順番に投げて計算時間を比較します．

compare.py

import time
import cpu_dense
import cpu_sparse
import gpu_dense

dt = 1.e-16
upto = 1.e-13
k = 1.e2
l = 1.e-10
m = 3.e-25
a = l*0.1

Nx = [5, 50, 50, 100, 100]
Ny = [3, 3, 30, 30, 50]

for nx, ny in zip(Nx, Ny):
    for main in [cpu_dense, cpu_sparse, gpu_dense]:
        start = time.perf_counter()
        wd = main.main(nx, ny, dt, upto, k, l, m, a)
        end = time.perf_counter()
        with open(wd/f'{end - start}.txt', 'w') as f:
            f.write(f'it took {end - start}.')

　そして計算時間を比較したグラフが下です．横軸が総原子数，縦軸は計算にかかった時間です．ゲーミングPCだけでなく手持ちのMac MiniでもCPU計算をさせてみました．

　**GPUの圧勝．**CPUだと2時間以上かかりそうなN=8000の条件でも6分で終わりました．~~インテル終わってる．~~CPU上での疎行列化も計算時間を半減させる琴葉できてますね．numpy-likeに簡単に扱えるので遊びで物理計算したい程度ならゲーミングPCのGPUでガンガン試せそうです．スパーステンソルが扱えないので変な工夫とか要らないのもメリットに思えてきます．
　じゃあ何でもGPUでいいかと言うと世の中そんなに甘くないです．グラフが途中で切れてるので察しがつくかと思いますが，GPUはVRAMの限界がN=8000で来てしまいました．2次元格子点の距離を保持させたテンソルの要素数は$2\times8000^2$です．これ以上のサイズで計算したければメモリに余裕のあるCPUを使わざるをえません．残念．

結論

　VRAMが溢れない限りはスパース化できてなくてもGPU使った方が絶対にいい．マジで泣きながらCPU使いましょう．

Continue...?

ブロック化してバッチ処理するかも