PyTorch显存管理：从分配机制到调试实战

显存管理，说白了就是跟GPU内存打交道。很多同学跑模型时遇到OOM（Out of Memory），第一反应就是“换更大的卡”。其实不然，很多时候是没搞清楚PyTorch怎么管显存的。今天我就带你彻底搞懂这一块。

一、PyTorch显存分配机制

PyTorch的显存分配，底层依赖的是CUDA的缓存分配器。我刚开始接触时也纳闷：为什么我明明只创建了一个小张量，显存却涨了一大截？

原因在于：PyTorch会预分配一块显存作为缓存。当你创建张量时，它先从缓存里拿；释放时，也不是真还给系统，而是放回缓存。这样做的目的是减少频繁的cudaMalloc/cudaFree调用——这两个操作很慢。

举个例子：

import torch

# 创建一个1GB的张量
x = torch.zeros(1024, 1024, 256, device='cuda')
print(torch.cuda.memory_allocated())  # 约1GB

# 删除它
del x
torch.cuda.empty_cache()
print(torch.cuda.memory_allocated())  # 接近0

# 但如果你不调用empty_cache()
x = torch.zeros(1024, 1024, 256, device='cuda')
del x
print(torch.cuda.memory_cached())  # 缓存还在，约1GB

嗯，这里要注意：empty_cache()会清空缓存，但不会减少已分配但未使用的显存。说白了，它只是把缓存还给系统，不是释放正在用的内存。

二、torch.cuda API详解

torch.cuda模块提供了丰富的显存管理API。我个人习惯把它们分成三类：查询类、控制类、调试类。

2.1 查询类API

API	作用	返回值
torch.cuda.memory_allocated()	当前已分配的显存（字节）	int
torch.cuda.memory_cached()	缓存中的显存（字节）	int
torch.cuda.memory_reserved()	保留的显存（字节），等于allocated + cached	int
torch.cuda.max_memory_allocated()	历史最大分配量	int
torch.cuda.memory_summary()	详细的显存使用报告	str

你想想看，如果想知道模型到底吃了多少显存，用memory_allocated最准。但如果你想看峰值，就得用max_memory_allocated。我在项目中遇到过一个问题：训练到一半OOM，但看当前分配量并不高。后来发现是某个中间变量在反向传播时膨胀了，峰值显存爆了。

2.2 控制类API

# 清空缓存
torch.cuda.empty_cache()

# 设置缓存分配器的参数
torch.cuda.set_per_process_memory_fraction(0.8)  # 限制使用80%显存

# 重置峰值统计
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()

set_per_process_memory_fraction这个API我强烈推荐。它就像给显存上了个“天花板”。我曾经在共享GPU的服务器上跑实验，就是因为没设这个限制，把别人的进程挤爆了。设成0.8，留20%给系统和其他进程，大家都舒服。

2.3 调试类API

# 开启内存统计
torch.cuda.memory._record_memory_history()

# 获取分配栈信息
snapshot = torch.cuda.memory._snapshot()
# 可以保存为pickle文件，用chrome的tracing工具查看

# 关闭统计
torch.cuda.memory._record_memory_history(None)

这个_record_memory_history是个隐藏神器。它记录了每次显存分配的调用栈。你想想看，当显存泄漏时，你能精确知道是哪个操作、哪行代码分配了显存没释放。我在优化一个百亿参数模型时，就是靠这个找到了一个循环里重复创建的临时张量。

三、显存监控与调试工具

光有API还不够，得会用工具。我常用的有三类：

3.1 nvidia-smi

最基础的监控工具。但要注意：它显示的是进程级别的显存占用，不是PyTorch内部的分配细节。

# 实时监控
watch -n 1 nvidia-smi

# 只看显存
nvidia-smi --query-gpu=memory.used,memory.free --format=csv

3.2 PyTorch Profiler

PyTorch自带的性能分析工具，也能看显存。

import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True
) as prof:
    # 你的训练代码
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()

print(prof.key_averages().table(sort_by="self_cuda_memory_usage"))

这个工具能告诉你每个操作消耗了多少显存。我记得有一次，发现一个看似无害的view操作竟然分配了大量显存——原来是因为它触发了张量的连续性检查，导致内存重新分配。

3.3 自定义监控装饰器

我习惯写一个简单的装饰器，在训练循环里自动记录显存变化：

import functools
import time

def monitor_memory(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        torch.cuda.reset_peak_memory_stats()
        start = time.time()
        
        result = func(*args, **kwargs)
        
        end = time.time()
        peak = torch.cuda.max_memory_allocated() / 1024**3
        current = torch.cuda.memory_allocated() / 1024**3
        
        print(f"[{func.__name__}] 耗时: {end-start:.2f}s, "
              f"峰值显存: {peak:.2f}GB, 当前显存: {current:.2f}GB")
        return result
    return wrapper

# 使用
@monitor_memory
def train_step(batch):
    # 训练逻辑
    pass

四、知识体系总览

下面这张图总结了本章的核心逻辑：

这张图把本章内容串起来了。从左到右，从底层机制到上层工具，层层递进。你掌握了分配机制，才能理解API为什么这么设计；理解了API，才能用好调试工具。

想要更多资料，联系可可