PyTorch训练后显存未释放？深度解析与高效清理指南

一、问题现象与开发者痛点

在PyTorch训练过程中，开发者常遇到这样的困惑：明明训练脚本已执行完毕，但通过nvidia-smi命令查看GPU显存占用时，仍显示大量显存被占用。这种”显存滞留”现象不仅浪费宝贵的GPU资源，更可能引发后续训练任务因显存不足而失败。特别是在多任务并行或云服务器环境中，显存管理不当会显著降低开发效率。

典型场景包括：

• 训练循环结束后，显存占用未降至初始水平
• 多次运行训练脚本后，可用显存逐渐减少
• 切换模型架构时，旧模型显存未完全释放
• 使用Jupyter Notebook时，内核重启后显存仍被占用

二、显存未释放的根源剖析

1. Python垃圾回收机制延迟

PyTorch张量对象受Python垃圾回收器(GC)管理，存在非确定性释放特性。当训练脚本结束时，若张量对象仍存在引用（如全局变量、闭包捕获等），GC不会立即回收这些对象。示例代码如下：

import torch
def train_model():
    model = torch.nn.Linear(1000, 1000).cuda()  # 全局变量未释放
    input_tensor = torch.randn(1000).cuda()
    output = model(input_tensor)
    return output
# 首次调用后显存被占用
_ = train_model()
# 此时显存可能未完全释放

2. CUDA上下文残留

PyTorch初始化时会创建CUDA上下文，该上下文会占用固定量的显存（约200-500MB）。即使所有张量被释放，此部分显存也不会自动释放。可通过以下代码验证：

import torch
print(torch.cuda.memory_allocated())  # 0
_ = torch.randn(1).cuda()  # 触发CUDA上下文创建
print(torch.cuda.memory_allocated())  # 非零值
torch.cuda.empty_cache()  # 仍无法释放上下文占用

3. 计算图保留

在训练过程中，若未正确断开计算图（如保留loss.backward()的中间结果），会导致整个计算图驻留内存。典型错误模式：

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
outputs = model(inputs)
loss = criterion(outputs, targets)
# 错误：保留计算图引用
grad_history = []
def record_grad():
    grad_history.append([p.grad for p in model.parameters()])
loss.register_hook(record_grad)  # 计算图被保留
loss.backward()  # 计算图无法释放

三、系统性解决方案

1. 显式内存管理策略

（1）手动清理缓存

import torch
def clear_gpu_memory():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()  # 清理缓存池
        # 强制Python垃圾回收
        import gc
        gc.collect()
        # 验证释放效果
        print(f"Allocated: {torch.cuda.memory_allocated()/1024**2:.2f}MB")
        print(f"Cached: {torch.cuda.memory_reserved()/1024**2:.2f}MB")
# 使用示例
model = torch.nn.Linear(1000, 1000).cuda()
_ = model(torch.randn(1000).cuda())
clear_gpu_memory()  # 显存占用显著下降

（2）上下文管理器模式

from contextlib import contextmanager
import torch
@contextmanager
def gpu_memory_manager():
    try:
        yield
    finally:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            import gc
            gc.collect()
# 使用示例
with gpu_memory_manager():
    model = torch.nn.Linear(1000, 1000).cuda()
    _ = model(torch.randn(1000).cuda())
# 退出with块后自动清理

2. 计算图优化技巧

（1）使用with torch.no_grad():

model.eval()
with torch.no_grad():  # 禁用梯度计算
    for inputs, targets in test_loader:
        outputs = model(inputs.cuda())
        # 推理代码...

（2）及时释放中间变量

# 错误模式：保留所有中间结果
outputs = model(inputs)
loss = criterion(outputs, targets)
# ...后续代码未使用outputs但仍保留
# 正确模式：显式删除无用变量
outputs = model(inputs)
loss = criterion(outputs, targets)
del outputs  # 立即释放

3. 高级调试方法

（1）显存分配追踪

def print_memory_usage():
    allocated = torch.cuda.memory_allocated()/1024**2
    reserved = torch.cuda.memory_reserved()/1024**2
    print(f"Allocated: {allocated:.2f}MB, Reserved: {reserved:.2f}MB")
# 在关键代码点插入检查
print_memory_usage()  # 初始状态
x = torch.randn(1000, 1000).cuda()
print_memory_usage()  # 分配后
del x
torch.cuda.empty_cache()
print_memory_usage()  # 清理后

（2）使用PyTorch Profiler

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    profile_memory=True
) as prof:
    # 测试代码
    x = torch.randn(1000, 1000).cuda()
    _ = x * 2
print(prof.key_averages().table(
    sort_by="cuda_memory_usage", row_limit=10))

四、最佳实践建议

1. 模块化设计：将模型训练封装为函数，避免全局变量污染
2. 资源释放顺序：遵循”数据→模型→优化器”的删除顺序
3. 监控常态化：在训练循环中定期打印显存使用情况
4. 异常处理：使用try-finally确保资源释放
5. 版本管理：保持PyTorch版本与CUDA驱动版本兼容

典型资源释放流程示例：

def train_and_cleanup():
    try:
        model = MyModel().cuda()
        optimizer = torch.optim.Adam(model.parameters())
        # 训练代码...
    finally:
        # 显式释放顺序
        del optimizer
        del model
        torch.cuda.empty_cache()
        import gc
        gc.collect()

五、特殊场景处理

1. 多GPU训练环境

在DataParallel或DistributedDataParallel模式下，需额外注意：

# DataParallel示例
model = torch.nn.DataParallel(MyModel()).cuda()
# 清理时需先解包
if isinstance(model, torch.nn.DataParallel):
    del model.module  # 先删除主模块
del model  # 再删除DP包装器

2. Jupyter Notebook环境

在Notebook中建议：

1. 使用%reset命令清除所有变量
2. 安装ipywidgets管理内核状态
3. 定期重启内核（Ctrl+M + .）

3. 云服务器环境

云GPU实例需特别注意：

• 设置自动回收策略（如AWS的Spot实例）
• 监控显存使用阈值（通过CloudWatch等）
• 实现训练任务超时自动终止机制

六、未来技术展望

PyTorch团队正在持续优化显存管理：

1. 即时编译(JIT)优化：减少中间变量存储
2. 统一内存管理：CPU-GPU内存自动交换
3. 更精细的垃圾回收：基于引用计数的即时释放

开发者可关注PyTorch GitHub仓库的#49312（显存管理优化）和#51208（计算图优化）等议题，及时获取最新进展。

结语

PyTorch显存管理是深度学习开发中的关键环节，需要开发者理解底层机制并掌握系统性的解决方案。通过显式清理策略、计算图优化和调试工具的综合运用，可有效解决训练后显存滞留问题。建议开发者建立标准化的资源管理流程，并结合监控工具实现自动化管理，从而提升开发效率和资源利用率。