Python 调用 CUDA 的准备工作

确保系统已安装 NVIDIA 显卡驱动、CUDA Toolkit 和 cuDNN 库。通过 nvidia-smi 命令检查驱动和 CUDA 版本是否兼容。安装 PyTorch 或 TensorFlow 的 GPU 版本时，需指定与 CUDA 版本匹配的安装命令。

验证 GPU 是否可用：

import torch
print(torch.cuda.is_available())  # 输出 True 表示可用

使用 PyTorch 实现 GPU 加速

将张量和模型转移到 GPU 上运行：

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 创建张量并转移到 GPU
x = torch.randn(1000, 1000).to(device)
y = torch.randn(1000, 1000).to(device)

# 自动利用 GPU 进行矩阵运算
z = torch.mm(x, y)

自定义 CUDA 核函数通过 PyTorch 的 torch.jit 实现：

from torch import jit

@jit.script
def gpu_operation(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a * b + a / (b + 1e-8)

使用 Numba 编写 CUDA 核函数

安装 Numba 库并编写直接运行的 GPU 核函数：

from numba import cuda
import numpy as np

@cuda.jit
def vector_add(a, b, result):
    idx = cuda.grid(1)
    if idx < len(result):
        result[idx] = a[idx] + b[idx]

# 生成数据并拷贝到 GPU
n = 100000
a = np.random.rand(n).astype(np.float32)
b = np.random.rand(n).astype(np.float32)
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_result = cuda.device_array_like(a)

# 配置线程块并执行核函数
threads_per_block = 256
blocks_per_grid = (n + threads_per_block - 1) // threads_per_block
vector_add[blocks_per_grid, threads_per_block](d_a, d_b, d_result)

# 将结果拷贝回主机
result = d_result.copy_to_host()

使用 CuPy 进行 GPU 数组运算

CuPy 提供与 NumPy 类似的接口，直接在 GPU 上执行运算：

import cupy as cp

# 创建 GPU 数组
x_gpu = cp.random.rand(10000)
y_gpu = cp.random.rand(10000)

# 自动在 GPU 上执行运算
z_gpu = x_gpu * y_gpu + cp.exp(x_gpu)

# 将结果传回 CPU
z_cpu = cp.asnumpy(z_gpu)

性能优化技巧

使用异步执行和流处理提高并行效率：

stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
    # 在此上下文中执行 GPU 操作
    output = model(input)

避免频繁的 CPU-GPU 数据传输，尽量在 GPU 上完成数据预处理。使用 torch.nn.DataParallel 或多进程实现多 GPU 并行计算。监控 GPU 使用情况：

watch -n 1 nvidia-smi