PyTorch学习阶段四-GPU调度 ： 资源管理与性能优化

PyTorch 架构级学习系列 - 第 5 篇

本文深入GPU调度的核心原理。你将理解CPU到GPU的数据传输、显存管理机制、性能瓶颈定位，以及混合精度训练等加速技术。

---

📚 目录

1. Device Context：CPU与GPU的数据流转
2. 显存管理：理解Caching Allocator
3. 性能优化：找出训练瓶颈
4. 混合精度训练：加速2倍的秘密
5. 常见OOM问题与解决方案

---

🔄 Part 1: Device Context —— CPU与GPU的数据流转

核心问题：数据在CPU和GPU之间如何移动？为什么会成为瓶颈？

1.1 理解设备（Device）

CPU vs GPU：

import torch

# CPU Tensor
cpu_tensor = torch.randn(1000, 1000)
print(cpu_tensor.device)  # cpu

# GPU Tensor
gpu_tensor = torch.randn(1000, 1000, device='cuda')
print(gpu_tensor.device)  # cuda:0

# 或者先创建再移动
cpu_tensor = torch.randn(1000, 1000)
gpu_tensor = cpu_tensor.to('cuda')  # CPU → GPU

三种指定设备的方式：

# 方式1：创建时指定
x = torch.randn(100, 100, device='cuda')

# 方式2：使用.to()
x = torch.randn(100, 100)
x = x.to('cuda')

# 方式3：使用.cuda()（不推荐，不够灵活）
x = torch.randn(100, 100)
x = x.cuda()

1.2 数据传输的成本

问题：CPU ↔ GPU 数据传输很慢！

import torch
import time

# 创建大tensor
size = 10000
cpu_a = torch.randn(size, size)
cpu_b = torch.randn(size, size)

# CPU上计算
start = time.time()
cpu_result = torch.matmul(cpu_a, cpu_b)
cpu_time = time.time() - start
print(f"CPU计算时间: {cpu_time:.4f}秒")

# GPU上计算（包含数据传输）
start = time.time()
gpu_a = cpu_a.to('cuda')      # CPU → GPU
gpu_b = cpu_b.to('cuda')      # CPU → GPU
gpu_result = torch.matmul(gpu_a, gpu_b)
torch.cuda.synchronize()      # 等待GPU计算完成
gpu_time = time.time() - start
print(f"GPU计算时间（含传输）: {gpu_time:.4f}秒")

# 纯GPU计算（数据已在GPU）
start = time.time()
gpu_result = torch.matmul(gpu_a, gpu_b)
torch.cuda.synchronize()
pure_gpu_time = time.time() - start
print(f"纯GPU计算时间: {pure_gpu_time:.4f}秒")

"""
典型输出：
CPU计算时间: 2.5秒
GPU计算时间（含传输）: 0.3秒
纯GPU计算时间: 0.01秒  ← 数据传输占了大部分时间！
"""

核心洞察：

• PCIe带宽瓶颈：CPU ↔ GPU 传输很慢（PCIe 3.0约16GB/s）
• GPU计算很快：一旦数据在GPU上，计算速度是CPU的100倍+
• 优化原则：尽量减少CPU ↔ GPU数据传输次数

1.3 异步执行与同步点

GPU操作是异步的：

import torch
import time

x = torch.randn(1000, 1000, device='cuda')

start = time.time()
y = torch.matmul(x, x)  # 立即返回，GPU在后台计算
print(f"matmul返回时间: {time.time() - start:.6f}秒")  # 几乎是0！

# 访问结果会触发同步
start = time.time()
result = y.cpu()  # 同步点：等待GPU计算完成
print(f"实际计算时间: {time.time() - start:.6f}秒")

常见同步点：

1. .cpu() - 数据从GPU返回CPU
2. .item() - 获取标量值
3. .numpy() - 转换为NumPy数组
4. torch.cuda.synchronize() - 显式同步
5. print(gpu_tensor) - 打印tensor

性能陷阱：

# ❌ 错误：频繁同步
losses = []
for epoch in range(100):
    loss = train_step()
    losses.append(loss.item())  # 每次都同步！很慢

# ✅ 正确：批量同步
loss_values = []
for epoch in range(100):
    loss = train_step()
    loss_values.append(loss)  # 先收集，不同步

losses = [l.item() for l in loss_values]  # 最后一起同步

1.4 模型的设备管理

移动整个模型：

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型
model = MyModel()

# 移动到GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 使用
x = torch.randn(32, 100).to(device)  # 输入也要在同一设备
output = model(x)

检查模型设备：

# 查看模型参数在哪个设备
for name, param in model.named_parameters():
    print(f"{name}: {param.device}")

# 检查所有参数是否在同一设备
def check_model_device(model):
    devices = {param.device for param in model.parameters()}
    if len(devices) > 1:
        print(f"⚠️ 警告：模型参数分布在多个设备: {devices}")
    else:
        print(f"✓ 模型在设备: {devices.pop()}")

---

小结：

• CPU ↔ GPU 数据传输是瓶颈（PCIe带宽限制）
• GPU操作是异步的，访问结果会触发同步
• 优化原则：让数据尽早到GPU，尽晚回CPU
• 模型和数据必须在同一设备

接下来，我们深入显存管理...

---

💾 Part 2: 显存管理 —— 理解Caching Allocator

核心问题：为什么nvidia-smi显示显存被占用，但实际没用那么多？

2.1 PyTorch的显存分配机制

Caching Allocator工作原理：

PyTorch请求显存 → 检查缓存池 → 有可用块？直接用 : 向CUDA申请新块
释放显存       → 不立即归还CUDA → 放入缓存池 → 等待复用

为什么这样设计？

• CUDA分配/释放显存很慢（毫秒级）
• 训练中频繁创建/销毁Tensor → 缓存避免频繁系统调用
• 代价：nvidia-smi显示的显存>实际使用

2.2 查看显存使用

三种查看方式：

import torch

# 1. torch内置方法（推荐）
print(f"已分配: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"已缓存: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

# 2. 详细统计
print(torch.cuda.memory_summary())

# 3. nvidia-smi（命令行）
# 显示的是 memory_reserved，不是 memory_allocated

示例：观察显存变化

import torch

def print_memory():
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"已分配: {allocated:.0f}MB, 已缓存: {reserved:.0f}MB")

print("初始状态:")
print_memory()

# 创建大tensor
x = torch.randn(10000, 10000, device='cuda')
print("\n创建10000x10000 tensor后:")
print_memory()  # 已分配: ~400MB, 已缓存: ~512MB

# 删除tensor
del x
print("\n删除tensor后:")
print_memory()  # 已分配: 0MB, 已缓存: 512MB（缓存未释放！）

# 清空缓存
torch.cuda.empty_cache()
print("\n清空缓存后:")
print_memory()  # 已分配: 0MB, 已缓存: 0MB

2.3 OOM（Out of Memory）诊断

症状： RuntimeError: CUDA out of memory

原因1：Batch Size太大

# 问题
batch_size = 256  # 太大
x = torch.randn(batch_size, 3, 224, 224, device='cuda')  # OOM!

# 解决：减小batch size
batch_size = 64  # 改小
x = torch.randn(batch_size, 3, 224, 224, device='cuda')  # OK

原因2：中间变量累积

# ❌ 问题：保留了计算图
losses = []
for epoch in range(1000):
    loss = model(x)
    losses.append(loss)  # loss是GPU tensor，保留了梯度！

# ✅ 解决：detach或只保存值
losses = []
for epoch in range(1000):
    loss = model(x)
    losses.append(loss.item())  # 只保存数值，释放GPU内存

原因3：没有清零梯度

# ❌ 问题
for epoch in range(100):
    loss = model(x)
    loss.backward()  # 梯度累积！
    optimizer.step()

# ✅ 解决
for epoch in range(100):
    optimizer.zero_grad()  # 清零梯度
    loss = model(x)
    loss.backward()
    optimizer.step()

2.4 显存优化技巧

技巧1：梯度累积（小batch模拟大batch）

# 问题：batch_size=256 OOM
# 解决：用4个batch_size=64累积梯度

accumulation_steps = 4
optimizer.zero_grad()

for i, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = criterion(output, target)
    loss = loss / accumulation_steps  # 缩放损失
    loss.backward()  # 累积梯度

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

技巧2：del + empty_cache

# 训练循环中释放大变量
for epoch in range(epochs):
    for data, target in train_loader:
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    # 验证阶段
    with torch.no_grad():
        val_output = model(val_data)  # 大变量

    # 立即释放
    del val_output
    torch.cuda.empty_cache()

技巧3：混合精度训练（下一节详述）

使用FP16可以减少显存占用50%！

---

小结：

• PyTorch使用Caching Allocator管理显存
• nvidia-smi显示的是缓存量，不是实际使用量
• OOM原因：batch太大、中间变量累积、忘记清零梯度
• 优化：梯度累积、及时删除变量、混合精度

接下来，我们学习如何定位性能瓶颈...

---

⚡ Part 3: 性能优化 —— 找出训练瓶颈

核心问题：训练慢是GPU慢？数据加载慢？还是其他原因？

3.1 性能分析工具

方法1：简单计时

import time

# 数据加载时间
data_time = 0
compute_time = 0

for i, (data, target) in enumerate(train_loader):
    start = time.time()
    data, target = data.to('cuda'), target.to('cuda')
    data_time += time.time() - start

    start = time.time()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    torch.cuda.synchronize()  # 等待GPU完成
    compute_time += time.time() - start

print(f"数据加载时间: {data_time:.2f}s")
print(f"计算时间: {compute_time:.2f}s")

方法2：PyTorch Profiler

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for i, (data, target) in enumerate(train_loader):
        if i >= 10:  # 只profile前10个batch
            break
        output = model(data.cuda())
        loss = criterion(output, target.cuda())
        loss.backward()
        optimizer.step()

# 打印统计
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

3.2 常见性能瓶颈

瓶颈1：数据加载慢

症状： GPU利用率低（<50%），大量时间在等待数据

解决方案：

# 1. 增加num_workers
train_loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,  # 多进程加载
    pin_memory=True  # 加速CPU→GPU传输
)

# 2. 使用预处理
# ❌ 在训练循环中处理
for data, target in train_loader:
    data = transform(data)  # 每次都要处理

# ✅ 在Dataset中预处理
class MyDataset(Dataset):
    def __getitem__(self, idx):
        data = self.data[idx]
        data = self.transform(data)  # 提前处理
        return data

瓶颈2：小batch导致GPU未充分利用

# GPU利用率低 → 增大batch size
batch_size = 32   # GPU利用率: 40%
batch_size = 128  # GPU利用率: 85%  ← 更好

瓶颈3：频繁的CPU-GPU同步

# ❌ 每步都同步
for i in range(1000):
    loss = train_step()
    print(f"Step {i}: {loss.item()}")  # 触发同步

# ✅ 批量打印
for i in range(1000):
    loss = train_step()
    if i % 100 == 0:  # 每100步打印一次
        print(f"Step {i}: {loss.item()}")

3.3 Benchmark：不同配置的性能对比

import torch
import time

def benchmark(model, input_size, batch_size, num_iterations=100):
    device = 'cuda'
    model = model.to(device)
    model.eval()

    # Warm up
    dummy_input = torch.randn(batch_size, *input_size, device=device)
    for _ in range(10):
        _ = model(dummy_input)

    # Benchmark
    torch.cuda.synchronize()
    start = time.time()

    for _ in range(num_iterations):
        _ = model(dummy_input)

    torch.cuda.synchronize()
    elapsed = time.time() - start

    throughput = (num_iterations * batch_size) / elapsed
    print(f"Batch size: {batch_size}, Throughput: {throughput:.2f} samples/sec")

# 测试不同batch size
model = MyModel()
for bs in [16, 32, 64, 128, 256]:
    try:
        benchmark(model, input_size=(3, 224, 224), batch_size=bs)
    except RuntimeError as e:
        print(f"Batch size {bs}: OOM")

---

小结：

• 使用Profiler定位瓶颈（数据加载？计算？）
• 数据加载慢：增加num_workers、预处理、pin_memory
• GPU利用率低：增大batch size
• 减少CPU-GPU同步次数

接下来，我们学习最强大的加速技术：混合精度训练...

---

🚀 Part 4: 混合精度训练 —— 加速2倍的秘密

核心思想：用FP16计算，但关键部分保持FP32精度

4.1 什么是混合精度？

精度对比：

FP32（全精度）：32位浮点数
- 范围: ±3.4 × 10³⁸
- 精度: 7位十进制
- 显存: 4 bytes/数

FP16（半精度）：16位浮点数
- 范围: ±65504
- 精度: 3位十进制
- 显存: 2 bytes/数  ← 显存减半！
- 速度: 2-3倍快（Tensor Core加速）

为什么不全用FP16？

• 数值范围小 → 容易溢出（overflow/underflow）
• 精度低 → 梯度太小会变成0

混合精度策略：

前向传播 → FP16（快）
损失计算 → FP32（精确）
反向传播 → FP16（快）+ 梯度缩放
参数更新 → FP32（精确）

4.2 使用PyTorch的AMP（Automatic Mixed Precision）

完整代码示例：

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

# 模型、优化器
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 梯度缩放器（防止FP16梯度下溢）
scaler = GradScaler()

# 训练循环
for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.cuda(), target.cuda()

        optimizer.zero_grad()

        # 自动混合精度上下文
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # 缩放损失并反向传播
        scaler.scale(loss).backward()

        # 更新参数
        scaler.step(optimizer)
        scaler.update()

关键API：

1. autocast() - 自动选择FP16/FP32
2. GradScaler() - 梯度缩放，防止下溢
3. scaler.scale(loss) - 放大损失
4. scaler.step() - 检查梯度，更新参数
5. scaler.update() - 调整缩放因子

4.3 梯度缩放原理

问题：FP16梯度下溢

# FP16最小正数: 6 × 10⁻⁵
# 如果梯度 = 1 × 10⁻⁷ → FP16中变成0！

# 解决：缩放
# 1. 放大损失 → 放大梯度
loss_scaled = loss * 1024  # 缩放因子
loss_scaled.backward()

# 2. 梯度自动被放大
# gradient = gradient * 1024

# 3. 更新参数前缩小回来
optimizer.step()  # scaler内部会除以1024

GradScaler动态调整缩放因子：

# 初始: scale = 65536
# 如果梯度正常 → scale *= 2（增大）
# 如果梯度inf/nan → scale /= 2（减小）

4.4 性能对比

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

def train_fp32(model, data, target, criterion, optimizer, num_iters=100):
    start = time.time()
    for _ in range(num_iters):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    torch.cuda.synchronize()
    return time.time() - start

def train_fp16(model, data, target, criterion, optimizer, num_iters=100):
    scaler = GradScaler()
    start = time.time()
    for _ in range(num_iters):
        optimizer.zero_grad()
        with autocast():
            output = model(data)
            loss = criterion(output, target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    torch.cuda.synchronize()
    return time.time() - start

# 测试
model = MyLargeModel().cuda()
data = torch.randn(64, 3, 224, 224).cuda()
target = torch.randint(0, 1000, (64,)).cuda()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

fp32_time = train_fp32(model, data, target, criterion, optimizer)
fp16_time = train_fp16(model, data, target, criterion, optimizer)

print(f"FP32: {fp32_time:.2f}s")
print(f"FP16: {fp16_time:.2f}s")
print(f"加速比: {fp32_time/fp16_time:.2f}x")

"""
典型输出：
FP32: 12.5s
FP16: 5.8s
加速比: 2.16x  ← 2倍加速！
"""

4.5 注意事项

何时使用混合精度？

✅ 适用：

• 大模型（参数>1M）
• 有Tensor Core的GPU（V100/A100/RTX 20/30系列）
• 训练任务（推理可以直接用FP16）

❌ 不适用：

• 小模型（加速不明显，甚至变慢）
• 老GPU（无Tensor Core）
• 对数值精度要求极高的任务

可能的问题：

# 问题1：精度损失导致不收敛
# 解决：调整学习率、增加epsilon

# 问题2：某些操作不支持FP16
# 解决：用autocast自动处理，或手动指定FP32
with autocast():
    # 大部分操作用FP16
    x = model(input)

    # 特定操作强制FP32
    with autocast(enabled=False):
        x = x.float()  # 转FP32
        x = some_sensitive_op(x)

---

小结：

• 混合精度：FP16计算+FP32精度 = 速度快+精度高
• AMP自动管理精度切换
• GradScaler防止梯度下溢
• 典型加速：2-3倍，显存减半
• 适用于大模型+新GPU

---

🔧 Part 5: 常见问题速查

5.1 OOM（显存不足）

原因排查：

1. Batch size太大 → 减小
2. 模型太大 → 使用梯度检查点（gradient checkpointing）
3. 中间变量未释放 → del + torch.cuda.empty_cache()
4. 梯度累积 → 使用混合精度

5.2 训练速度慢

排查步骤：

1. GPU利用率低？ → 增大batch size、检查数据加载
2. 数据加载慢？ → 增加num_workers、pin_memory
3. 频繁同步？ → 减少.item()调用
4. 小模型？ → 使用混合精度

5.3 精度问题

混合精度导致不收敛：

• 调整学习率（可能需要稍微增大）
• 某些层强制FP32
• 检查数值是否溢出

---

🎓 总结

核心要点

Device Management:

• CPU ↔ GPU传输是瓶颈
• 让数据尽早到GPU，尽晚回CPU
• GPU操作是异步的

显存管理:

• Caching Allocator会缓存显存
• memory_allocated vs memory_reserved
• OOM：减小batch、梯度累积、混合精度

性能优化:

• Profiler定位瓶颈
• 数据加载：num_workers、pin_memory
• 减少同步点

混合精度:

• 2-3倍加速，显存减半
• AMP自动管理
• GradScaler防止梯度下溢

快速参考

# 设备管理
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
data = data.to(device)

# 显存查看
print(f"已分配: {torch.cuda.memory_allocated()/1024**3:.2f}GB")
print(f"已缓存: {torch.cuda.memory_reserved()/1024**3:.2f}GB")

# 混合精度
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast():
    output = model(data)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 梯度累积
accumulation_steps = 4
for i, (data, target) in enumerate(train_loader):
    loss = model(data)
    loss = loss / accumulation_steps
    loss.backward()
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()