CuPy vs Numba vs PyTorch：GPU 加速方案怎么选

# CuPy：按 CUDA 版本选包，如 CUDA 12.x
pip install cupy-cuda12x

# Numba：自带 LLVM，CPU/GPU 均可
pip install numba

# PyTorch（CUDA 12.1 示例）
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

import numpy as np, cupy as cp, torch, numba
from numba import cuda
print("numpy", np.__version__)
print("cupy ", cp.__version__, "| cuda", cp.cuda.runtime.runtimeGetVersion())
print("torch", torch.__version__, "| cuda available:", torch.cuda.is_available())

import numpy as np
# ---- CPU 版（原 Numpy 代码）----
x_cpu = np.random.rand(8000, 8000)
%time y_cpu = x_cpu @ x_cpu          # 约数百 ms

# ---- 仅改 import，逻辑不变 ----
import cupy as cp
x_gpu = cp.random.rand(8000, 8000)   # 分配在显存
y_gpu = x_gpu @ x_gpu                # 在 GPU 执行
out   = cp.asnumpy(y_gpu)            # 取回主机（必要时才搬）

import cupy as cp

def bench_cupy(n=8000):
    x = cp.random.rand(n, n)
    y = cp.random.rand(n, n)
    cp.cuda.Stream.null.synchronize()        # 等之前的任务清空
    start = cp.cuda.Event()
    end   = cp.cuda.Event()
    start.record()
    z = x @ y
    end.record()
    end.synchronize()                        # 阻塞到 kernel 完成
    ms = cp.cuda.get_elapsed_time(start, end)  # 真实 GPU 耗时（毫秒）
    return ms

print(f"matmul {bench_cupy():.2f} ms")

import cupy as cp, torch

x_cp = cp.random.rand(1000, 1000)
# cupy -> torch，无需 .get() 回主机
x_th = torch.from_dlpack(x_cp)              # 共享显存，零拷贝
print(x_th.device)                          # cuda:0
# torch -> cupy
x_back = cp.from_dlpack(x_th)

from numba import njit
import numpy as np

@njit
def pairwise_cpu(a, b, out):
    # 纯 Python 循环被编译成机器码
    for i in range(a.shape[0]):
        s = 0.0
        for j in range(a.shape[1]):
            s += a[i, j] * b[j, i]
        out[i] = s

a = np.random.rand(4000, 4000)
b = np.random.rand(4000, 4000)
out = np.empty(4000)
pairwise_cpu(a, b, out)        # 首次有编译开销，之后接近 C 速度

from numba import cuda
import numpy as np

@cuda.jit
def matvec_kernel(A, x, out):
    # 每个线程负责 out 的一行
    i = cuda.grid(1)                  # 全局一维索引
    if i < out.shape[0]:
        s = 0.0
        for j in range(A.shape[1]):
            s += A[i, j] * x[j]
        out[i] = s

A = np.random.rand(1 << 16, 512)
x = np.random.rand(512)
out = np.empty(A.shape[0])

d_A, d_x, d_out = cuda.to_device(A), cuda.to_device(x), cuda.to_device(out)
threads = 256
blocks  = (A.shape[0] + threads - 1) // threads   # 向上取整铺满
matvec_kernel[blocks, threads](d_A, d_x, d_out)
result = d_out.copy_to_host()                     # 结果搬回主机

from numba import vectorize

@vectorize(['float64(float64, float64)'], target='cuda')
def gpu_mul(a, b):
    return a * b

a = np.random.rand(10_000_000)
b = np.random.rand(10_000_000)
c = gpu_mul(a, b)        # 自动在 GPU 上逐元素并行

import torch

x = torch.randn(8000, 8000, device='cuda', requires_grad=True)
y = torch.randn(8000, 8000, device='cuda')
z = (x @ y).sum()
z.backward()                       # 自动求梯度，x.grad 已就绪
print(x.grad.shape)                # (8000, 8000)

import torch
from torch.cuda.amp import autocast, GradScaler

model = torch.nn.Linear(4096, 4096).cuda()
opt = torch.optim.SGD(model.parameters(), lr=1e-3)
scaler = GradScaler()

for step in range(100):
    inp = torch.randn(256, 4096, device='cuda')
    opt.zero_grad()
    with autocast():                      # 前向自动转 float16
        loss = model(inp).pow(2).mean()
    scaler.scale(loss).backward()         # 缩放梯度，防 underflow
    scaler.step(opt)
    scaler.update()

import torch
class Net(torch.nn.Module):
 def forward(self, x):
 return x * 2 + 1
ep = torch.export.export(Net(), (torch.randn(4, 4, device='cuda'),))
print(ep.graph_module) # 可序列化，便于部署

import cupy as cp, torch
from torch.utils.data import TensorDataset, DataLoader
# 1) CuPy 做向量化特征工程（留在显存）
raw = cp.random.rand(1_000_000, 64)
feat = (raw - raw.mean(axis=0)) / raw.std(axis=0) # 标准化
feat = cp.ascontiguousarray(feat)
# 2) DLPack 零拷贝交给 PyTorch
x = torch.from_dlpack(feat) # 不回主机
y = torch.randn(1_000_000, 1, device='cuda')
loader = DataLoader(TensorDataset(x, y), batch_size=4096, shuffle=True)
net = torch.nn.Sequential(torch.nn.Linear(64, 1)).cuda()
opt = torch.optim.Adam(net.parameters())
for xb, yb in loader:
 opt.zero_grad()
 loss = torch.nn.functional.mse_loss(net(xb), yb)
 loss.backward(); opt.step()

def recommend(use_case: str) -> str:
 return {
 "已有 numpy 代码想提速":   "CuPy（改 import）",
 "自定义循环/分支算法":      "Numba（@njit 或 @cuda.jit）",
 "深度学习/需求导":          "PyTorch",
 "预处理+训练混合":          "CuPy 预处理 + PyTorch 训练",
 }.get(use_case, "先量化瓶颈在带宽还是算力")