如何使用 numba 加速 np.corrcoef

Question

我的代码很简单：

@njit()
def corr(arr: np.ndarray):
    return np.corrcoef(arr)

arr = np.random.random([10000, 10000])
corr_matrix = corr(arr)

在我的电脑上大约需要 50 秒才能完成，没有

@njit

只需 18 秒。如果我将大小增加到 30,000，该函数将永远运行。

有没有办法提高 numba

@njit

在

np.corrcoef

上的性能，就像在这种情况下，或者

np.corrcoef

已经尽可能快了？我想我在这里没有正确理解 numba，因为使用

@njit

比不使用 numba 慢很多。

Answer 1

我在 Colab T4 GPU 实例上进行了测试。

from time import perf_counter_ns
import cupy as cp
import numpy as np

rng = np.random.default_rng(65651651684)
x = rng.random((3000, 3000))

%timeit np.corrcoef(x)
# 835 ms ± 14.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit cp.asnumpy(cp.corrcoef(x))
# 273 ms ± 2.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit cp.asnumpy(cp.corrcoef(x, dtype=cp.float32))
# 37.8 ms ± 581 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

如果我们担心缓存或延迟评估的可能性，我们可以生成新数据，运行一次，甚至包括一个

print

语句以进行良好的衡量。

rng = np.random.default_rng(23368551581)
x = rng.random((3000, 3000))

tic = perf_counter_ns()
rho0 = np.corrcoef(x)
print(rho0[0])
toc = perf_counter_ns()
print(toc - tic)  # 879117900

tic = perf_counter_ns()
rho = cp.asnumpy(cp.corrcoef(x))
print(rho[0])
toc = perf_counter_ns()
print(toc - tic)  # 445459866

# passes
np.testing.assert_allclose(cp.asnumpy(rho), rho0)

tic = perf_counter_ns()
rho = cp.asnumpy(cp.corrcoef(x, dtype=cp.float32))
print(rho[0])
toc = perf_counter_ns()
print(toc - tic)  # 41231956

# passes
np.testing.assert_allclose(cp.asnumpy(rho), rho0, atol=1e-6)