一、PyTorch显存占用问题的本质与影响

PyTorch作为深度学习框架的核心，其显存管理机制直接影响模型训练效率。显存占用过高会导致程序崩溃、训练中断，甚至引发多任务并行时的资源冲突。显存占用的主要来源包括模型参数（weights/biases）、中间计算结果（activations）、梯度（gradients）和优化器状态（optimizer states）。例如，一个包含1亿参数的模型，仅参数本身就可能占用400MB显存（FP32精度），若加上梯度则翻倍至800MB。

显存泄漏的典型场景包括：未释放的临时张量、循环中累积的计算图、未正确释放的CUDA上下文。例如，以下代码会导致显存持续占用：

import torch
for _ in range(100):
    x = torch.randn(1000, 1000).cuda()  # 每次循环创建新张量但未释放
    y = x @ x  # 计算结果未被回收

二、PyTorch显存清空的核心方法

1. 显式释放张量资源

通过del语句和torch.cuda.empty_cache()组合实现彻底释放：

import torch
# 创建占用显存的张量
x = torch.randn(10000, 10000).cuda()
y = x.clone()
# 显式释放
del x, y  # 删除Python对象引用
torch.cuda.empty_cache()  # 清空CUDA缓存池

原理：del仅删除Python对象引用，而empty_cache()会触发CUDA的内存管理器回收未使用的显存块。

2. 上下文管理器控制显存

自定义上下文管理器实现训练阶段的显存隔离：

from contextlib import contextmanager
@contextmanager
def clear_cuda_cache():
    try:
        yield
    finally:
        torch.cuda.empty_cache()
# 使用示例
with clear_cuda_cache():
    model = torch.nn.Linear(1000, 1000).cuda()
    input = torch.randn(64, 1000).cuda()
    output = model(input)

3. 梯度清零与优化器重置

在训练循环中，需区分zero_grad()和显存释放：

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(10):
    optimizer.zero_grad()  # 清零梯度但不释放显存
    output = model(input)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    # 强制释放计算图
    if epoch % 5 == 0:
        del output, loss
        torch.cuda.empty_cache()

三、显存占用监控与诊断工具

1. 内置工具nvidia-smi

终端实时监控命令：

watch -n 1 nvidia-smi -l 1  # 每秒刷新一次

输出字段解析：

• Used/Total：当前使用量/总显存
• GPU-Util：计算单元利用率
• Memory-Usage：显存占用百分比

2. PyTorch内置诊断

# 获取当前显存分配
print(torch.cuda.memory_allocated())  # 当前Python进程占用的显存
print(torch.cuda.max_memory_allocated())  # 历史峰值
# 详细分配记录（需启用跟踪）
torch.cuda.reset_peak_memory_stats()  # 重置统计
# 执行某些操作后...
print(torch.cuda.max_memory_reserved())  # 缓存池保留量

3. 第三方工具py3nvml

安装与使用：

pip install py3nvml

from py3nvml.py3nvml import *
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print(f"总显存: {info.total/1024**2:.2f}MB")
print(f"已用显存: {info.used/1024**2:.2f}MB")
nvmlShutdown()

四、显存优化高级策略

1. 混合精度训练

通过torch.cuda.amp减少显存占用：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

效果：FP16存储可减少50%显存占用，同时保持数值稳定性。

2. 梯度检查点（Gradient Checkpointing）

牺牲计算时间换取显存：

from torch.utils.checkpoint import checkpoint
def custom_forward(x):
    # 原始前向传播
    return model(x)
# 使用检查点
input = torch.randn(64, 1000).cuda()
output = checkpoint(custom_forward, input)

原理：仅存储输入输出而非中间激活，显存占用可降低至O(√N)。

3. 模型并行与张量并行

对于超大模型（如GPT-3），采用分片策略：

# 示例：参数分片到两个GPU
model_part1 = ModelPart1().cuda(0)
model_part2 = ModelPart2().cuda(1)
# 前向传播时同步
with torch.cuda.device(0):
    output1 = model_part1(input)
with torch.cuda.device(1):
    output2 = model_part2(output1)

五、常见问题解决方案

1. CUDA Out of Memory错误处理

try:
    output = model(input)
except RuntimeError as e:
    if "CUDA out of memory" in str(e):
        torch.cuda.empty_cache()
        # 尝试减小batch size或使用梯度累积
        small_input = input[:32]  # 减小batch
        output = model(small_input)
    else:
        raise

2. 多进程训练显存隔离

使用torch.multiprocessing时显式指定设备：

def train_worker(rank, world_size):
    torch.cuda.set_device(rank)
    # 每个进程独立管理显存
    model = Model().cuda(rank)
    ...
if __name__ == "__main__":
    mp.spawn(train_worker, args=(world_size,), nprocs=world_size)

3. 持久化缓存管理

通过环境变量控制缓存行为：

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"
# 限制每次分配的最大块大小，减少碎片

六、最佳实践总结

1. 监控常态化：在训练循环中定期记录显存使用情况
2. 释放及时化：对临时变量采用del+empty_cache()组合
3. 精度混合化：对非敏感层采用FP16
4. 检查点启用：对长序列模型默认开启
5. 碎片预防：设置合理的分配策略（如max_split_size_mb）

通过系统化的显存管理，可使PyTorch训练效率提升30%-50%，尤其在资源受限的环境下效果显著。开发者应根据具体场景选择组合策略，平衡计算速度与显存占用。