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

引言：数据进得来，才算得动

大规模数据采集常被目标站点限流、封 IP、验证码拦截。企业级代理 IP 服务如 亿牛云可让采集在分布式节点间平滑切换出口，保障数据稳定入库。数据进来后，瓶颈转向算力——下面用代码直接对比 CuPy、Numba、PyTorch 三种上 GPU 的方式。

0. 环境准备

# 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())

1. CuPy：把 Numpy 代码搬上 GPU

1.1 逐行迁移

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)            # 取回主机（必要时才搬）

1.2 正确的 GPU 计时（别用 time.time）

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")

1.3 显存与 DLPack 零拷贝互转

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)

2. Numba：JIT 编译与手写 CUDA 内核

2.1 CPU 加速：@njit

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 速度

2.2 GPU 内核：@cuda.jit（手动排布 grid/block）

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()                     # 结果搬回主机

2.3 声明式并行：@vectorize（免手写内核）

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 上逐元素并行

3. PyTorch：张量 + 自动微分 + AMP

3.1 张量上 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)

3.2 混合精度训练（AMP）压显存提速度

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()

3.3 导出推理图（脱离训练环境）

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)            # 可序列化，便于部署

4. 三者串联：CuPy 预处理 + PyTorch 训练

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()

5. 决策速查

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

选型铁律：用 cupy.cuda.Event / torch.cuda.Event 测真实 GPU 耗时，首跑丢弃（冷启动），并把 cpu→gpu 拷贝计入总账——小数据下 GPU 可能更慢。