numpyのexp, log等はtorchで高速化できる

numpyのexp, log等を使用していましたが，かなりの回数繰り返し呼び出しをする必要があり無視できないほど時間がかかったのでその際に発見した知見に関して書きます．

結論から言うとtorchから呼び出してnumpyに戻すという方法が比較的きれいなコードを保ちつつ高速に実行できます．

試した方法

numexpr

１つ目の方法はnumexprを利用する方法です．
numexprではevaluate関数を用いてnumpyを実行することによって（なぜか）高速化できるものとなっています．
ただし，scipyには対応していません．
この方法はCPU上の動作としては最速ですが，evaluateを利用する必要があるため，flake8のignoreを利用する必要性やerror messageがわかりづらくなってしまうなどの問題点があります．

torch

2つ目の方法はtorchを利用する方法です．
torch側からはnumpyを受け付けていないこと，torch.Tensorが少し特殊なことからinput, output時にそれぞれ変換が必要になります．
後述の実験からもわかりますが，torch.as_tensorの方が配列のcopyを行わない分高速です．
また，cudaを利用する場合は単純なexpだけであればさらになる高速化が見込めます（CUDAによる高速化のスニペットを見てください）が，deviceにデータを送る動作を考慮すると十分高速とは言えません．
つまり，cuda programmingをする場合はcpu -- cuda間の往来を最小に抑えない限りは高速化に繋がらない点に注意が必要です．

CPUからGPUに送る際に発生するLatency

rnd = np.random.random((10000, 10000))
%timeit rnd_cuda = torch.as_tensor(rnd).to("cuda")
73.8 ms ± 262 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Benchmarkingではcudaは利用しておりませんが，単純なexp処理だけなら約400倍早くなることを確認しました．(NVIDIA GeForce RTX 3060, CUDA11.4, Nvidia driver=470.103.01)

CUDAによる高速化

shape = (10000, 10000)
rnd_numpy = np.random.random(shape)
rnd_cuda = torch.rand(shape).to(shape)

%timeit np.exp(rnd_numpy)
>>> 968 ms ± 16.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit torch.exp(rnd_cuda)
>>> 2.56 ms ± 58.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Benchmark test

以下，これらのmodulesのimportと変数，関数の定義を前提とします．

以下のTestsで前提となるコード

import numexpr as ne  # 2.7.1
import numpy as np  # 1.20.1
import torch  # 1.11.0
from scipy.special import erf  # scipy==1.7.3
from scipy.special import logsumexp

shape = (10000, 10000)
rnd = np.random.random(shape)

def to_numpy(x: torch.Tensor) -> np.ndarray:
    return x.cpu().detach().numpy()

また，CPUはCore i7 10700Fを利用しました．（Core i7 10510U搭載のLaptopで試した際にはnumexprの方が高速だったのでお使いの環境で試されると良いと思います．）

1つ目のtestはexpです．

expのRuntime tests

%timeit np.exp(rnd)
>>> 965 ms ± 5.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit ne.evaluate("exp(rnd)")
>>> 89.8 ms ± 1.37 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.exp(torch.as_tensor(rnd)))
>>> 88.8 ms ± 1.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.exp(torch.tensor(rnd)))
>>> 168 ms ± 4.04 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

2つ目はlogです．

logのRuntime tests

%timeit np.log(rnd)
>>> 1.77 s ± 15.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit ne.evaluate("log(rnd)")
>>> 95.6 ms ± 6.52 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.log(torch.as_tensor(rnd)))
>>> 81.9 ms ± 2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.log(torch.tensor(rnd)))
>>> 163 ms ± 1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

次にerfを試します．(numpyにerfが存在しないため，numexprは近似値を計算できるtanhを実行しています)

erf(正規分布の積分に使われる誤差関数)のRuntime tests

%timeit erf(rnd)
>>> 387 ms ± 3.02 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit ne.evaluate("tanh(rnd)")
>>> 89.1 ms ± 1.41 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.erf(torch.as_tensor(rnd)))
>>> 93.6 ms ± 1.68 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit to_numpy(torch.erf(torch.tensor(rnd)))
>>> 171 ms ± 2.09 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

最後にlogsumexpを試します．(numpyにlogsumexpが存在しないこと・numexprは呼び出し内部でのnestができないこと，これらの理由からnumexprでは（厳密な実装とは違いますが）階層的に各moduleを呼び出す関数を利用します．)

logsumexpのRuntime tests

%timeit logsumexp(rnd, axis=0)
>>> 1.2 s ± 13.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

def ne_logsumexp(x):
    x = ne.evaluate("exp(x)")
    x = ne.evaluate("sum(x, axis=0)")
    return ne.evaluate("log(x)")

%timeit ne_logsumexp(rnd)
>>> 282 ms ± 3.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit to_numpy(torch.logsumexp(torch.as_tensor(rnd), axis=0))
>>> 240 ms ± 2.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit to_numpy(torch.logsumexp(torch.tensor(rnd), axis=0))
>>> 312 ms ± 5.75 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

結論

1. torchに存在しないnumpy, scipyのmethodsもありますが，それらに依存する必要がない場合はtorchを利用するのが良いです．
2. torchに依存したくない場合であっても，np.ndarray --> torch methods --> np.ndarrayという操作で十分にnumpyよりも高速なので，簡単に高速化したい場合はwrapperを作りましょう．
3. numexprは環境によってはCPU上でtorchよりも早く動く可能性があるので，（安全性度外視で）ダメ押しの高速化をどうしても行いたいときのみ使いましょう．