4. Python GPU 生态：RAPIDS、PyTorch/TensorFlow、Numba CUDA

好了，讲完了 CUDA 编程的理论，咱们得回到 Python 阵地。说实话，真正做数据分析的人，谁愿意天天写 C++ 和 CUDA C？我刚开始接触 GPU 加速时，也想过「我直接调 PyTorch 不就行了吗？」。但后来发现，纯 PyTorch 搞 DataFrame 操作太别扭了。于是 RAPIDS 出现了——它把 GPU 加速带到了 pandas 和 scikit-learn 的世界。

这一章，我们聊三个东西：RAPIDS 全家桶、深度学习框架的 GPU 基础、Numba 的 @cuda.jit 魔法。嗯，你想想看，搞多组学数据处理，哪一步能离开 DataFrame 操作和机器学习？所以这章的核心就是——让你只用 Python，就能把 GPU 用起来。

本章路线图：从数据分析层（cuDF） → 机器学习层（cuML） → 图分析层（cuGraph） → 深度学习（PyTorch/TF） → 自定义 CUDA kernel（Numba）。

4.1 RAPIDS 全家桶：cuDF、cuML、cuGraph

RAPIDS 是 NVIDIA 开源的 GPU 加速数据分析平台。它的 API 几乎完全模仿 pandas、scikit-learn、NetworkX，但底层跑在 GPU 上。我个人习惯在单机多组学数据处理时，直接用 cuDF 替代 pandas——只要数据能塞进显存，速度能快 10~50 倍。

4.1.1 cuDF：GPU 上的 DataFrame

cuDF 的接口和 pandas 高度兼容。下面是个简单例子：

import cudf
import pandas as pd

# 从 CSV 加载到 GPU
gdf = cudf.read_csv('rnaseq_counts.csv')
# 过滤和分组
result = gdf[gdf['gene'] != 'MT-'] \
           .groupby('sample')['count'] \
           .mean()
print(result)

是不是很像 pandas？注意一点：cuDF 的 DataFrame 默认在 GPU 显存里。在项目中，我踩过一个坑——直接 gdf.to_pandas() 把整个表拷回 CPU，结果内存炸了。所以我建议，尽量在 GPU 上完成全部操作，最后只导出小结果。

性能技巧：cuDF 的 GroupBy、Join、Sort 比 pandas 快很多，但一次性插入大量行不如分批 Bulk Load。用 cudf.concat() 合并多个小块数据，速度最优。

4.1.2 cuML：GPU 上的机器学习

cuML 提供了和 scikit-learn 一模一样的 API。你只需要把训练数据放在 cuDF DataFrame 或 cuPy 数组里，然后调 fit/predict 就行。

from cuml import RandomForestClassifier
from cuml.preprocessing import StandardScaler

X = gdf[['feature1','feature2','feature3']]
y = gdf['label']

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

model = RandomForestClassifier(n_estimators=100)
model.fit(X_scaled, y)

我曾经用 cuML 处理过单细胞 RNA-seq 的聚类任务，100 万细胞 × 2000 个高变基因，scikit-learn 跑了 6 分钟，cuML 只用 12 秒——你想想看，差别有多大。

但要小心：cuML 目前部分算法（如 t-SNE、UMAP）的随机种子行为可能和 CPU 版本不完全一致，用于生产前务必验证。

4.1.3 cuGraph：GPU 上的图分析

多组学中经常要构建基因-基因共表达网络、蛋白质互作网络等。cuGraph 用 GPU 加速了 PageRank、Louvain 社区发现、最短路径等算法。

import cugraph

# 从 cuDF 边表创建图
edge_gdf = cudf.DataFrame({
    'src': [0, 1, 2],
    'dst': [1, 2, 0],
    'weight': [0.5, 1.0, 0.8]
})
G = cugraph.from_cudf_edgelist(edge_gdf, source='src', destination='dst', weight='weight')

# Louvain 社区检测
parts, mod = cugraph.louvain(G)
print(parts.head())
注意：cuGraph 的图数据要求顶点 ID 是连续的整数。如果原始 ID 是字符串或跳跃整数，必须先用 cugraph.renumber() 重新编号。我曾经因为忘了这步，调试了整整一下午。

4.2 PyTorch/TensorFlow GPU 加速基础

搞多组学深度学习（比如从图像切片预测基因表达、用 Transformer 做单细胞表示学习），框架的 GPU 加速是基本操作。这里我只讲最关键的几点——说实话，这两年框架封装得太好了，配好环境后基本不用手动管理 GPU。

4.2.1 PyTorch GPU 设置

import torch

# 检查 GPU 是否可用
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')

# 将模型和数据移到 GPU
model = MyModel().to(device)
data = torch.randn(64, 3, 224, 224).to(device)
output = model(data)

# 多 GPU 数据并行
model = torch.nn.DataParallel(model)

我个人的习惯是：先用 torch.cuda.is_available() 做 fallback 处理，否则在无 GPU 的测试环境会直接报错。另外，注意只把需要计算的数据移到 GPU，不要一股脑把整个 Dataset 全放进去——显存满了就等着 OOM 吧。

4.2.2 TensorFlow GPU 设置

import tensorflow as tf

# 列出物理 GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    # 需要设置内存增长，避免一次占满
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

# 指定使用的 GPU（单机多卡时很有用）
tf.config.set_visible_devices(gpus[0], 'GPU')
经验：TensorFlow 默认会一次性占用所有显存。在多人共享 GPU 的服务器上，你肯定不想被骂。所以 set_memory_growth 或者手动 set_soft_device_placement 一定要加上。

两者对比一下：

特性PyTorchTensorFlow 2.x
GPU 设备管理手动 .to(device)自动 placement + 配置
混合精度torch.cuda.amptf.keras.mixed_precision
多卡并行DataParallel / DistributedDataParalleltf.distribute.MirroredStrategy
内存控制torch.cuda.empty_cache()tf.config.experimental.set_memory_growth

4.3 Numba CUDA：在 Python 里写 GPU 内核

有时候，cuDF 和 PyTorch 都帮不了你。比如你要实现一个自定义的基因计数合并算法，用纯 Python 循环会慢到哭。这个时候，Numba 就派上用场了——它让你用 Python 语法写 CUDA kernel，然后 JIT 编译到 GPU 上运行。

4.3.1 基本用法

from numba import cuda
import numpy as np

@cuda.jit
def vector_add(a, b, c):
    idx = cuda.grid(1)  # 一维线程索引
    if idx < a.size:
        c[idx] = a[idx] + b[idx]

# 准备数据
n = 1024
a = np.random.randn(n).astype(np.float32)
b = np.random.randn(n).astype(np.float32)
c = np.zeros_like(a)

# 将数据拷贝到 GPU
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_c = cuda.to_device(c)

# 配置线程块
threads_per_block = 256
blocks_per_grid = (n + threads_per_block - 1) // threads_per_block

# 启动 kernel
vector_add[blocks_per_grid, threads_per_block](d_a, d_b, d_c)

# 结果拷回 CPU
d_c.copy_to_host(c)

嗯，这里要注意：Numba 的 @cuda.jit 编译速度快但限制也不少——不支持 Python 的异常、动态类型、部分高级特性。我在项目中用它写过单细胞 UMI 去重，比纯 Python 快了 200 倍，但调试起来也比较痛苦，经常要打印中间值。

4.3.2 与 RAPIDS 搭配

你可以把 cuDF 的列转成 numpy 数组传给 Numba kernel，处理完再转回 cuDF。这样做的好处是：既享受了 cuDF 的数据管理便利，又能跑自定义并行逻辑。

实战组合拳：cuDF 加载数据 → 过滤/排序 → 取列转为 numpy → Numba CUDA kernel 复杂计算 → 结果写回 cuDF → cuML 训练。

4.4 本章知识结构

Python GPU 生态结构 用户 Python 脚本 RAPIDS 全家桶 PyTorch / TensorFlow Numba CUDA cuDF / cuML / cuGraph GPU 加速训练 / 推理 @cuda.jit / 自定义 kernel CUDA Runtime & GPU 硬件 从高层 API 到底层 kernel，Python GPU 生态覆盖了多组学数据分析的所有层面 组合使用 cuDF+cuML+Numba 能解决 90% 的 GPU 加速需求

4.5 避坑与建议

  • 显存不够怎么办？ —— 用 cuDF 的 cudf.concat() 分批加载，或转到 PyTorch DataLoader 做内存映射。我曾经处理过 100GB 的 scATAC-seq 数据，根本不可能全放显存，后来用 cuDF 的 chunked read 解决了。
  • cuML 结果和 sklearn 不一致？ —— 常见原因：浮点运算顺序、随机种子、算法实现细节。建议先在小样本上对比验证，差异在 1e-6 以内可接受。
  • Numba 的 CUDA 调试 —— 尽量用 @cuda.jit(debug=True) 编译，但会降低性能。我一般先在 CPU 上写纯 Python 版本验证逻辑，再改成 kernel。
  • 混合精度 —— PyTorch 用 torch.cuda.amp.autocast()，TF 用 tf.keras.mixed_precision.set_global_policy('mixed_float16')，显存减半，速度提升 1.5~2 倍。
重要提醒：cuda.to_device() 和 copy_to_host() 是同步操作，会阻塞 CPU 线程。如果在一系列 kernel 调用之间频繁拷贝数据，性能会急剧下降。尽量把数据留在 GPU 上，直到最后再取回。

好了，这一章的内容就到这儿。Python GPU 生态的工具链已经非常成熟了——你写代码时基本感觉不到 GPU 的存在，但速度却快了几个数量级。好好利用这些工具，多组学数据预处理不再是瓶颈。

← 上一章 📖 返回目录 下一章 →