NVIDIA CUDAを使って、NumPyを動かします。
> conda install -c conda-forge cupy
テスト
5000×5000行列の積を求めます。
import time
import numpy as np
import cupy as cp
# CPU
start_time = time.time()
a = np.random.rand(5000, 5000)
b = np.random.rand(5000, 5000)
result = np.dot(a, b)
end_time = time.time()
print(f"NumPy Time: {end_time - start_time} seconds")
# GPU
start_time = time.time()
a = cp.random.rand(5000, 5000)
b = cp.random.rand(5000, 5000)
result = cp.dot(a, b)
end_time = time.time()
print(f"CuPy Time: {end_time - start_time} seconds")
【結果】
NumPy Time: 0.8619685173034668 seconds
CuPy Time: 0.2240922451019287 seconds
CuPyの方が3.8倍速いかな。
