PyTorch GPU / CUDA 加速
深度学习的核心操作是大规模矩阵乘法与元素运算。CPU 的设计目标是处理复杂的串行逻辑，核心数通常为 8~64 个；而 GPU 拥有数千个简单并行核心，天然适合这类高度并行的数值计算。PyTorch 通过 NVIDIA 的 CUDA（Compute Unified Device Architecture）框架调用 GPU，可将训练速度提升数十倍乃至百倍。
1. CPU 与 GPU 的差异
GPU 训练速度的提升主要来自两方面：其一是并行执行大量相同计算；其二是高带宽显存使得数据搬运速度远快于 CPU 内存。
对于矩阵乘法这类运算密集型操作，加速效果尤为显著。
| 对比项 | CPU | GPU（NVIDIA） |
|---|---|---|
| 核心数量 | 8~64 个大核心 | 数千个小核心（CUDA Cores） |
| 设计目标 | 低延迟串行处理 | 高吞吐并行计算 |
| 内存带宽 | ~50~100 GB/s | ~500~3000 GB/s |
| 矩阵乘法速度 | 基准 | 快 10x~100x |
| PyTorch 接口 | "cpu" |
"cuda" |
如果使用 Apple Silicon（M1/M2/M3），PyTorch 通过 mps 后端支持 Metal GPU 加速，用法与 CUDA 几乎完全相同，只需将设备改为 "mps"。
2. 检测 CUDA 环境
使用 GPU 前，需要先确认当前环境是否安装了支持 CUDA 的 PyTorch。
同时需要检查系统上是否存在可用的 GPU 设备。
实例
import torch
# 是否支持 CUDA
print("CUDA 可用:", torch.cuda.is_available())
# GPU 设备数量
print("GPU 数量:", torch.cuda.device_count())
# 当前默认 GPU 的索引
print("当前 GPU:", torch.cuda.current_device())
# GPU 型号名称
print("GPU 型号:", torch.cuda.get_device_name(0))
# PyTorch 版本与编译时绑定的 CUDA 版本
print("PyTorch 版本:", torch.__version__)
print("CUDA 版本:", torch.version.cuda)
# 是否支持 CUDA
print("CUDA 可用:", torch.cuda.is_available())
# GPU 设备数量
print("GPU 数量:", torch.cuda.device_count())
# 当前默认 GPU 的索引
print("当前 GPU:", torch.cuda.current_device())
# GPU 型号名称
print("GPU 型号:", torch.cuda.get_device_name(0))
# PyTorch 版本与编译时绑定的 CUDA 版本
print("PyTorch 版本:", torch.__version__)
print("CUDA 版本:", torch.version.cuda)
输出示例：
CUDA 可用: True GPU 数量: 1 当前 GPU: 0 GPU 型号: NVIDIA GeForce RTX 4090 PyTorch 版本: 2.3.0+cu121 CUDA 版本: 12.1
2.1 动态选择设备（推荐写法）
在代码中硬编码 "cuda" 会导致没有 GPU 的机器直接报错。
实例
import torch
# 方式一：经典写法，兼容性最好
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 方式二：PyTorch 2.0+ 推荐，自动支持 CUDA / MPS / CPU
device = (
"cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu"
)
print(f"使用设备: {device}")
# 方式一：经典写法，兼容性最好
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 方式二：PyTorch 2.0+ 推荐，自动支持 CUDA / MPS / CPU
device = (
"cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu"
)
print(f"使用设备: {device}")
3. 张量在设备间移动
PyTorch 中的张量默认在 CPU 上创建。
要使用 GPU 计算，需要显式将张量移动到 GPU，或直接在 GPU 上创建。
3.1 基本移动方法
实例
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 创建 CPU 张量
cpu_tensor = torch.tensor([1.0, 2.0, 3.0])
print(cpu_tensor.device) # cpu
# 方式一：.to(device) —— 推荐，通用性最好
gpu_tensor = cpu_tensor.to(device)
# 方式二：.cuda() —— 仅限 CUDA 环境
gpu_tensor = cpu_tensor.cuda()
# 方式三：创建时直接指定设备
gpu_tensor = torch.tensor([1.0, 2.0, 3.0], device=device)
gpu_tensor = torch.randn(3, 4, device=device)
# 移回 CPU（用于打印、numpy 转换、保存等操作）
back_to_cpu = gpu_tensor.cpu()
print(gpu_tensor.device) # cuda:0
print(back_to_cpu.device) # cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 创建 CPU 张量
cpu_tensor = torch.tensor([1.0, 2.0, 3.0])
print(cpu_tensor.device) # cpu
# 方式一：.to(device) —— 推荐，通用性最好
gpu_tensor = cpu_tensor.to(device)
# 方式二：.cuda() —— 仅限 CUDA 环境
gpu_tensor = cpu_tensor.cuda()
# 方式三：创建时直接指定设备
gpu_tensor = torch.tensor([1.0, 2.0, 3.0], device=device)
gpu_tensor = torch.randn(3, 4, device=device)
# 移回 CPU（用于打印、numpy 转换、保存等操作）
back_to_cpu = gpu_tensor.cpu()
print(gpu_tensor.device) # cuda:0
print(back_to_cpu.device) # cpu
GPU 张量转 numpy 时，必须先移回 CPU，且若张量附带梯度还需先 detach：
实例
# 普通 GPU 张量转 numpy
arr = gpu_tensor.cpu().numpy()
# 带梯度的 GPU 张量转 numpy
arr = gpu_tensor.detach().cpu().numpy()
arr = gpu_tensor.cpu().numpy()
# 带梯度的 GPU 张量转 numpy
arr = gpu_tensor.detach().cpu().numpy()
3.2 设备一致性约束
不同设备上的张量不能直接参与同一运算。
否则会抛出 RuntimeError：
实例
a = torch.randn(3).to("cuda")
b = torch.randn(3) # 在 CPU 上
# c = a + b # RuntimeError: Expected all tensors to be on the same device
# 正确做法：先统一设备
c = a + b.to("cuda")
b = torch.randn(3) # 在 CPU 上
# c = a + b # RuntimeError: Expected all tensors to be on the same device
# 正确做法：先统一设备
c = a + b.to("cuda")
查看张量所在设备：
实例
x = torch.randn(3, 4).to(device)
print(x.device) # cuda:0
print(x.is_cuda) # True
print(x.get_device()) # 0（GPU 索引）
print(x.device) # cuda:0
print(x.is_cuda) # True
print(x.get_device()) # 0（GPU 索引）
3.3 速度对比验证
实例
import torch
import time
device = torch.device("cuda")
n = 5000
# CPU 矩阵乘法
a_cpu = torch.randn(n, n)
b_cpu = torch.randn(n, n)
start = time.time()
c_cpu = torch.matmul(a_cpu, b_cpu)
print(f"CPU 耗时: {time.time() - start:.3f}s")
# GPU 矩阵乘法
a_gpu = a_cpu.to(device)
b_gpu = b_cpu.to(device)
torch.cuda.synchronize() # 确保数据已传输完毕再开始计时
start = time.time()
c_gpu = torch.matmul(a_gpu, b_gpu)
torch.cuda.synchronize() # 等待 GPU 执行完毕再停止计时
print(f"GPU 耗时: {time.time() - start:.3f}s")
import time
device = torch.device("cuda")
n = 5000
# CPU 矩阵乘法
a_cpu = torch.randn(n, n)
b_cpu = torch.randn(n, n)
start = time.time()
c_cpu = torch.matmul(a_cpu, b_cpu)
print(f"CPU 耗时: {time.time() - start:.3f}s")
# GPU 矩阵乘法
a_gpu = a_cpu.to(device)
b_gpu = b_cpu.to(device)
torch.cuda.synchronize() # 确保数据已传输完毕再开始计时
start = time.time()
c_gpu = torch.matmul(a_gpu, b_gpu)
torch.cuda.synchronize() # 等待 GPU 执行完毕再停止计时
print(f"GPU 耗时: {time.time() - start:.3f}s")
输出示例：
CPU 耗时: 1.847s GPU 耗时: 0.021s
GPU 计算是异步执行的——Python 调用返回后，GPU 上的运算可能尚未完成。计时时必须调用
torch.cuda.synchronize()等待 GPU 真正结束，否则测量结果不准确。
4. 模型移动到 GPU
模型的所有参数（weight、bias）本质上也是张量。
这些参数同样需要移动到 GPU 才能在 GPU 上执行前向传播和反向传播。
对整个模型调用 .to(device) 即可，PyTorch 会自动遍历并移动所有内部参数。
实例
import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 10),
)
def forward(self, x):
return self.net(x)
# 模型移动到 GPU，只需调用一次
model = SimpleNet().to(device)
# 验证参数是否都在 GPU 上
for name, param in model.named_parameters():
print(f"{name}: {param.device}")
# net.0.weight: cuda:0
# net.0.bias: cuda:0
# ...
# 输入数据也必须在相同设备上
x = torch.randn(32, 784).to(device)
output = model(x)
print(output.shape) # torch.Size([32, 10])
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 10),
)
def forward(self, x):
return self.net(x)
# 模型移动到 GPU，只需调用一次
model = SimpleNet().to(device)
# 验证参数是否都在 GPU 上
for name, param in model.named_parameters():
print(f"{name}: {param.device}")
# net.0.weight: cuda:0
# net.0.bias: cuda:0
# ...
# 输入数据也必须在相同设备上
x = torch.randn(32, 784).to(device)
output = model(x)
print(output.shape) # torch.Size([32, 10])
模型在 GPU 上，但输入数据仍在 CPU 上时，前向传播会报错。务必在 DataLoader 读取每个 batch 之后，将
inputs和labels都调用.to(device)。
5. 完整训练流程
以下是一个包含数据加载、模型训练、验证评估的标准 GPU 训练模板。
实例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 1. 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"训练设备: {device}")
# 2. 数据加载
# pin_memory=True：将数据锁定在内存，加快 CPU -> GPU 传输
# num_workers：多进程预加载，减少数据等待时间
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
train_dataset = datasets.MNIST(root="./data", train=True,
download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False,
download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True,
num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False,
num_workers=4, pin_memory=True)
# 3. 定义模型
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 7 * 7, 256), nn.ReLU(), nn.Dropout(0.5),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
model = CNN().to(device) # 模型移到 GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 4. 训练函数
def train_epoch(model, loader, optimizer, criterion):
model.train()
total_loss, correct = 0.0, 0
for inputs, labels in loader:
# non_blocking=True：异步传输，CPU 可以继续准备下一批数据
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 5. 验证函数
def eval_epoch(model, loader, criterion):
model.eval()
total_loss, correct = 0.0, 0
with torch.no_grad():
for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
outputs = model(inputs)
loss = criterion(outputs, labels)
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 6. 主训练循环
for epoch in range(1, 11):
train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion)
val_loss, val_acc = eval_epoch(model, test_loader, criterion)
print(f"Epoch {epoch:02d} | "
f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 1. 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"训练设备: {device}")
# 2. 数据加载
# pin_memory=True：将数据锁定在内存，加快 CPU -> GPU 传输
# num_workers：多进程预加载，减少数据等待时间
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
train_dataset = datasets.MNIST(root="./data", train=True,
download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False,
download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True,
num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False,
num_workers=4, pin_memory=True)
# 3. 定义模型
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 7 * 7, 256), nn.ReLU(), nn.Dropout(0.5),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
model = CNN().to(device) # 模型移到 GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 4. 训练函数
def train_epoch(model, loader, optimizer, criterion):
model.train()
total_loss, correct = 0.0, 0
for inputs, labels in loader:
# non_blocking=True：异步传输，CPU 可以继续准备下一批数据
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 5. 验证函数
def eval_epoch(model, loader, criterion):
model.eval()
total_loss, correct = 0.0, 0
with torch.no_grad():
for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
outputs = model(inputs)
loss = criterion(outputs, labels)
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 6. 主训练循环
for epoch in range(1, 11):
train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion)
val_loss, val_acc = eval_epoch(model, test_loader, criterion)
print(f"Epoch {epoch:02d} | "
f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")
关键要点：
device变量在代码最顶部声明，全局统一引用，不硬编码"cuda"- 模型
.to(device)只需调用一次 - 每个 batch 的
inputs和labels都要.to(device) - DataLoader 设置
pin_memory=True和non_blocking=True加速数据传输 - 验证阶段使用
torch.no_grad()关闭梯度，节省显存
6. 多 GPU 训练
当单卡显存不足时，可以考虑使用多 GPU 并行训练。
此外，多 GPU 还能进一步提升训练速度。PyTorch 提供了两种主要方式。
6.1 DataParallel
这是最简单的多卡方式。
它采用单进程，将每个 batch 均分到各 GPU，前向传播并行执行，梯度在主卡上汇总更新。使用方便，但主卡负担重，多卡利用率不均衡，适合快速入门。
实例
import torch
import torch.nn as nn
model = CNN()
if torch.cuda.device_count() > 1:
print(f"使用 {torch.cuda.device_count()} 个 GPU")
model = nn.DataParallel(model)
# 也可以指定使用哪些 GPU
# model = nn.DataParallel(model, device_ids=[0, 1])
model = model.to("cuda")
# 之后的训练代码与单卡完全一致
# DataParallel 会自动将 batch 均分到各 GPU，并汇总结果
import torch.nn as nn
model = CNN()
if torch.cuda.device_count() > 1:
print(f"使用 {torch.cuda.device_count()} 个 GPU")
model = nn.DataParallel(model)
# 也可以指定使用哪些 GPU
# model = nn.DataParallel(model, device_ids=[0, 1])
model = model.to("cuda")
# 之后的训练代码与单卡完全一致
# DataParallel 会自动将 batch 均分到各 GPU，并汇总结果
如果需要访问原始模型的属性（如自定义方法），需要通过 .module 访问：
实例
# model 被 DataParallel 包裹后，原始模型在 model.module 中
print(model.module.classifier)
# 保存模型时建议保存 model.module，方便单卡加载
torch.save(model.module.state_dict(), "model.pth")
print(model.module.classifier)
# 保存模型时建议保存 model.module，方便单卡加载
torch.save(model.module.state_dict(), "model.pth")
6.2 DistributedDataParallel
这是多进程方式。
每个进程绑定一块 GPU，每张卡持有完整模型，通过 All-Reduce 同步梯度。是生产环境和大规模训练的推荐方案，效率远高于 DataParallel。
实例
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def main(rank, world_size):
# 初始化进程组，nccl 后端专为 NVIDIA GPU 优化
dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
# 每个进程绑定对应的 GPU
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")
# 模型移到对应 GPU 后包装 DDP
model = CNN().to(device)
model = DDP(model, device_ids=[rank])
# DataLoader 使用 DistributedSampler，保证各卡数据不重叠
sampler = DistributedSampler(train_dataset,
num_replicas=world_size, rank=rank)
loader = DataLoader(train_dataset, batch_size=64,
sampler=sampler, pin_memory=True)
# 训练逻辑与单卡完全一致
# ...
dist.destroy_process_group()
# 启动方式（推荐使用 torchrun）：
# torchrun --nproc_per_node=4 train_ddp.py
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def main(rank, world_size):
# 初始化进程组，nccl 后端专为 NVIDIA GPU 优化
dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
# 每个进程绑定对应的 GPU
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")
# 模型移到对应 GPU 后包装 DDP
model = CNN().to(device)
model = DDP(model, device_ids=[rank])
# DataLoader 使用 DistributedSampler，保证各卡数据不重叠
sampler = DistributedSampler(train_dataset,
num_replicas=world_size, rank=rank)
loader = DataLoader(train_dataset, batch_size=64,
sampler=sampler, pin_memory=True)
# 训练逻辑与单卡完全一致
# ...
dist.destroy_process_group()
# 启动方式（推荐使用 torchrun）：
# torchrun --nproc_per_node=4 train_ddp.py
| 对比项 | DataParallel | DistributedDataParallel |
|---|---|---|
| 进程数量 | 单进程 | 多进程（每卡一个） |
| 通信后端 | Python GIL 限制 | NCCL（高效） |
| 主卡负担 | 重（汇总梯度） | 均衡（All-Reduce） |
| 代码改动 | 极少 | 中等 |
| 适用场景 | 快速实验 | 生产训练 |
7. 混合精度训练 AMP
默认训练使用 float32（FP32）精度。
自动混合精度（Automatic Mixed Precision，AMP）让部分计算使用 float16（FP16）或 bfloat16（BF16），在几乎不损失精度的情况下获得显著收益：
- 显存占用减少约 50%
- 训练速度提升 2x~3x（Tensor Core 硬件加速）
- 代码改动极少，只需三处修改
实例
from torch.cuda.amp import autocast, GradScaler
model = CNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler() # FP16 梯度缩放器，防止梯度下溢出为零
for epoch in range(num_epochs):
model.train()
for inputs, labels in train_loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
# 修改一：autocast 区域内自动选择 FP16/FP32
with autocast(device_type="cuda"):
outputs = model(inputs)
loss = criterion(outputs, labels)
# 修改二：用 scaler 缩放 loss 再反向传播
scaler.scale(loss).backward()
# 修改三：scaler 更新参数，内部自动处理梯度缩放
scaler.step(optimizer)
scaler.update()
model = CNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler() # FP16 梯度缩放器，防止梯度下溢出为零
for epoch in range(num_epochs):
model.train()
for inputs, labels in train_loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
# 修改一：autocast 区域内自动选择 FP16/FP32
with autocast(device_type="cuda"):
outputs = model(inputs)
loss = criterion(outputs, labels)
# 修改二：用 scaler 缩放 loss 再反向传播
scaler.scale(loss).backward()
# 修改三：scaler 更新参数，内部自动处理梯度缩放
scaler.step(optimizer)
scaler.update()
7.1 FP16 与 BF16 的选择
| 格式 | 精度位 | 数值范围位 | 适合硬件 | 特点 |
|---|---|---|---|---|
| float16 | 10 位 | 5 位 | RTX 20/30/40 系 | 需要 GradScaler 防溢出 |
| bfloat16 | 7 位 | 8 位 | A100 / H100 / RTX 4090 | 范围与 FP32 相同，训练更稳定 |
如果 GPU 支持 BF16，优先使用，无需 GradScaler：
实例
# BF16 写法（PyTorch 2.0+）
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
8. 性能优化技巧
本节介绍多种优化显存和训练速度的技巧。
8.1 显存管理
实例
# 查看当前显存使用情况
print(f"已分配: {torch.cuda.memory_allocated() / 1024**2:.1f} MB")
print(f"已缓存: {torch.cuda.memory_reserved() / 1024**2:.1f} MB")
# 打印详细显存报告
print(torch.cuda.memory_summary())
# 释放缓存池（不释放已分配的显存）
torch.cuda.empty_cache()
# 推理时使用 inference_mode（比 no_grad 更快，彻底禁用梯度引擎）
with torch.inference_mode():
output = model(x)
# 梯度检查点：以重新计算换显存（大模型训练常用）
from torch.utils.checkpoint import checkpoint
output = checkpoint(model_block, x)
print(f"已分配: {torch.cuda.memory_allocated() / 1024**2:.1f} MB")
print(f"已缓存: {torch.cuda.memory_reserved() / 1024**2:.1f} MB")
# 打印详细显存报告
print(torch.cuda.memory_summary())
# 释放缓存池（不释放已分配的显存）
torch.cuda.empty_cache()
# 推理时使用 inference_mode（比 no_grad 更快，彻底禁用梯度引擎）
with torch.inference_mode():
output = model(x)
# 梯度检查点：以重新计算换显存（大模型训练常用）
from torch.utils.checkpoint import checkpoint
output = checkpoint(model_block, x)
8.2 DataLoader 优化
实例
loader = DataLoader(
dataset,
batch_size=256,
num_workers=4, # 多进程预加载，建议设为 CPU 核数的 50%
pin_memory=True, # 锁页内存，加快 CPU -> GPU 传输
persistent_workers=True, # 进程常驻，避免每个 epoch 重新创建进程
prefetch_factor=2, # 每个 worker 预取的 batch 数量
)
dataset,
batch_size=256,
num_workers=4, # 多进程预加载，建议设为 CPU 核数的 50%
pin_memory=True, # 锁页内存，加快 CPU -> GPU 传输
persistent_workers=True, # 进程常驻，避免每个 epoch 重新创建进程
prefetch_factor=2, # 每个 worker 预取的 batch 数量
)
8.3 torch.compile 模型编译（PyTorch 2.0+）
一行代码，对计算图进行编译优化，速度提升 30%~200%（取决于模型结构）：
实例
# 编译模型，第一次运行时有编译开销，之后的迭代速度显著提升
model = torch.compile(model)
# 不同编译模式的权衡
model = torch.compile(model, mode="default") # 均衡，适合大多数场景
model = torch.compile(model, mode="reduce-overhead") # 降低调度开销
model = torch.compile(model, mode="max-autotune") # 最大优化，编译时间较长
model = torch.compile(model)
# 不同编译模式的权衡
model = torch.compile(model, mode="default") # 均衡，适合大多数场景
model = torch.compile(model, mode="reduce-overhead") # 降低调度开销
model = torch.compile(model, mode="max-autotune") # 最大优化，编译时间较长
8.4 梯度累积（模拟大 batch）
显存不足时，可以通过梯度累积模拟更大的 batch size，而无需实际增大显存占用：
实例
accumulation_steps = 8 # 每 8 个 batch 更新一次参数，等效 batch_size × 8
for i, (inputs, labels) in enumerate(train_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels) / accumulation_steps # 均分 loss
loss.backward() # 梯度累积，不清零
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # 只在参数更新后清零
for i, (inputs, labels) in enumerate(train_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels) / accumulation_steps # 均分 loss
loss.backward() # 梯度累积，不清零
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # 只在参数更新后清零
8.5 显存不足时的应对策略
| 策略 | 方法 | 效果 |
|---|---|---|
| 减小 batch size | 256 -> 64 | 线性减少显存 |
| 混合精度 | AMP + FP16/BF16 | 显存减少约 50% |
| 梯度累积 | 每 N 步更新一次 | 模拟大 batch，不增加显存 |
| 梯度检查点 | torch.utils.checkpoint |
大幅减少显存，训练速度降低 |
| 冻结部分参数 | 迁移学习冻结主干 | 减少反向传播的显存占用 |
9. 常见错误与排查
本节汇总了 GPU 训练中最常见的错误及其解决方案。
9.1 RuntimeError: Expected all tensors to be on the same device
原因：参与运算的张量分布在不同设备上（一个在 CPU，一个在 GPU）。
实例
# 排查方式：打印各张量的 device
print(inputs.device, labels.device, next(model.parameters()).device)
# 解决：确保每个 batch 都执行了 .to(device)
inputs = inputs.to(device)
labels = labels.to(device)
print(inputs.device, labels.device, next(model.parameters()).device)
# 解决：确保每个 batch 都执行了 .to(device)
inputs = inputs.to(device)
labels = labels.to(device)
9.2 RuntimeError: CUDA out of memory
原因：显存不足，常见于 batch size 过大、模型过大，或计算图意外积累。
实例
# 常见隐患：循环中 loss 没有调用 .item()，导致计算图不断堆积
# 错误写法
total_loss += loss # loss 是张量，持有整个计算图
# 正确写法
total_loss += loss.item() # .item() 取出 Python 标量，释放计算图
# 其他排查步骤：
# 1. 减小 batch_size
# 2. 确认验证循环使用了 torch.no_grad()
# 3. 调用 torch.cuda.empty_cache() 清理缓存
# 4. 使用 torch.cuda.memory_summary() 定位显存占用来源
# 错误写法
total_loss += loss # loss 是张量，持有整个计算图
# 正确写法
total_loss += loss.item() # .item() 取出 Python 标量，释放计算图
# 其他排查步骤：
# 1. 减小 batch_size
# 2. 确认验证循环使用了 torch.no_grad()
# 3. 调用 torch.cuda.empty_cache() 清理缓存
# 4. 使用 torch.cuda.memory_summary() 定位显存占用来源
9.3 Can't call numpy() on Tensor that requires grad
原因：带梯度的张量不能直接转为 numpy 数组。
实例
# 错误写法
arr = gpu_tensor.numpy()
# 正确写法
arr = gpu_tensor.detach().cpu().numpy()
arr = gpu_tensor.numpy()
# 正确写法
arr = gpu_tensor.detach().cpu().numpy()
9.4 CUDA error: device-side assert triggered
原因：通常是标签值越界（如 10 分类但 label 值等于 10），或数组索引越界。报错信息在 GPU 上产生，默认显示位置不准确。
实例
# 设置环境变量，让错误在准确的代码行处抛出（会变为同步执行，速度变慢）
import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
9.5 训练速度没有提升（GPU 利用率低）
原因：数据加载成为瓶颈，GPU 大部分时间在等待 CPU 准备数据。
排查步骤：
- 运行
nvidia-smi或watch -n 1 nvidia-smi观察 GPU 利用率- 利用率持续低于 80% 说明存在数据瓶颈
- 增大 DataLoader 的 num_workers
- 开启 pin_memory=True
- 将数据预处理移到 GPU（torchvision.transforms 支持 GPU 操作）
- 将小文件数据集预先加载到内存
9.6 使用 PyTorch Profiler 精确定位瓶颈
实例
from torch.profiler import profile, ProfilerActivity
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
) as prof:
for i, (inputs, labels) in enumerate(train_loader):
if i >= 10:
break
inputs = inputs.to(device)
labels = labels.to(device)
loss = criterion(model(inputs), labels)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# 按 CUDA 耗时排序，打印前 15 项
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
) as prof:
for i, (inputs, labels) in enumerate(train_loader):
if i >= 10:
break
inputs = inputs.to(device)
labels = labels.to(device)
loss = criterion(model(inputs), labels)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# 按 CUDA 耗时排序，打印前 15 项
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
10. API 快速参考
以下是常用 PyTorch GPU 相关 API 的速查表。
10.1 设备管理
| 操作 | 代码 |
|---|---|
| 检测 CUDA 可用 | torch.cuda.is_available() |
| 选择设备 | torch.device("cuda" if ... else "cpu") |
| GPU 数量 | torch.cuda.device_count() |
| GPU 名称 | torch.cuda.get_device_name(0) |
| 设置当前 GPU | torch.cuda.set_device(0) |
| 等待 GPU 完成 | torch.cuda.synchronize() |
10.2 张量操作
| 操作 | 代码 |
|---|---|
| 张量移到 GPU | tensor.to(device) 或 tensor.cuda() |
| 张量移到 CPU | tensor.cpu() |
| 查看所在设备 | tensor.device |
| 是否在 GPU | tensor.is_cuda |
| GPU 张量转 numpy | tensor.detach().cpu().numpy() |
| 异步传输 | tensor.to(device, non_blocking=True) |
10.3 模型操作
| 操作 | 代码 |
|---|---|
| 模型移到 GPU | model.to(device) |
| 开启训练模式 | model.train() |
| 开启推理模式 | model.eval() |
| 编译加速 | torch.compile(model) |
| 关闭梯度 | with torch.no_grad(): |
| 推理模式 | with torch.inference_mode(): |
10.4 显存管理
| 操作 | 代码 |
|---|---|
| 已分配显存 | torch.cuda.memory_allocated() |
| 已缓存显存 | torch.cuda.memory_reserved() |
| 显存报告 | torch.cuda.memory_summary() |
| 清理缓存 | torch.cuda.empty_cache() |
10.5 核心原则速记
1. 用 device 变量统一管理，不要硬编码 "cuda" 2. 模型和数据必须在同一设备上，每个 batch 都要 .to(device) 3. 验证和推理时一定使用 torch.no_grad() 或 torch.inference_mode() 4. 生产训练推荐开启 AMP 混合精度，几乎免费获得 2x 加速 5. DataLoader 设置 pin_memory=True 和 num_workers >= 4 减少数据瓶颈 6. PyTorch 2.0+ 可以用 torch.compile(model) 一行提速 30%+ 7. GPU 计算是异步的，精确计时需要调用 torch.cuda.synchronize()