【Python】從 NumPy 到 CuPy：運用 CUDA 實現 GPU 加速矩陣運算

CuPy 與 NumPy 的介紹與差異

NumPy 是在 CPU 上運行的數值計算庫，對於大部分中小型數據運算已足夠，但在大規模運算中效能有限；而 CuPy 基於 NVIDIA CUDA，與 NumPy API 高度相似，讓使用者能夠輕鬆地將原本在 CPU 上運算的程式轉換到 GPU 上執行，從而加速大規模運算任務。

安裝 CuPy

請先在終端機輸入 nvcc --version 檢查你應該安裝甚麼版本的 CuPy。如果有顯示 Cuda compilation tools, release 12.1, V12.1.105 代表你應該安裝 cupy-cuda12x。

【注意】如果程式使用 CuPy 的途中會導致藍屏，代表你應該更新一下顯示卡的驅動程式。

nvcc --version // 檢查你應該安裝甚麼版本的 CuPy
pip install cupy-cuda12x // 安裝指定版本 CuPy

展示 CuPy 的強大功能

為了生動展示 CuPy 的強大功能，可以選擇一個既簡單又能引起興趣的目標：矩陣相乘的 GPU 加速。

我們可以先使用 NumPy 在 CPU 實現矩陣相乘，隨後再使用 CuPy 在 GPU 實現矩陣相乘，觀察兩者的差異。

為什麼我要選擇示範 「 矩陣相乘」 ?

先說明一下，矩陣相乘通常是一個耗時的運算，尤其當矩陣的規模較大時。如果有兩個矩陣都是 n×n 的方陣，那麼計算它們的乘法通常需要 O(n³) 的時間，這使得該操作在規模較大的情況下非常耗費時間。正因為如此，這個範例能夠清楚地展示 CuPy 在平行處理大量數據時的優勢。

import numpy as np
import time

# 生成 2 個 1000x1000 的隨機矩陣
A = np.random.rand(1000, 1000)
B = np.random.rand(1000, 1000)

start_time = time.time() # 計算開始時間

# 求矩陣相乘 1000 次
for i in range(1000):
    np.dot(A, B)

end_time = time.time() # 計算結束時間
t = end_time - start_time # 計算耗時

print(f"CPU 計算矩陣相乘 1000 次，耗時 : {t:.10f} 秒")

CPU 計算矩陣相乘 1000 次，耗時 : 20.9240744114 秒

import cupy as cp
import time

# 以防萬一，等待 GPU 上所有目前的任務都完成後 ( 例如釋放記憶體 )，才繼續執行後續程式碼。
cp.cuda.Device(0).synchronize()

# 生成 2 個 1000x1000 的隨機矩陣
A = cp.random.rand(1000, 1000)
B = cp.random.rand(1000, 1000)

start_time = time.time() # 計算開始時間

# 求矩陣相乘 1000 次
for i in range(1000):
    cp.dot(A, B)

end_time = time.time() # 計算結束時間
t = end_time - start_time # 計算耗時

print(f"GPU 計算矩陣相乘 1000 次，耗時 : {t:.10f} 秒")

GPU 計算矩陣相乘 1000 次，耗時 : 0.0220127106 秒

結語

在本次測試中，我們發現使用 CuPy（GPU）進行 2 個 1000 × 1000 矩陣相乘，重複 1000 次的運算，與 NumPy（CPU）相比，速度快了約 950 倍（結果可能會因為硬體差異而有所不同）！這突顯了 GPU 在大規模線性代數運算中的強大加速效能。