第2章：PyTorch显存管理机制

显存管理，说白了就是PyTorch怎么跟GPU打交道。我刚开始做模型训练时，经常遇到OOM（显存溢出），那时候我以为是模型太大，后来才发现——很多时候是我不懂PyTorch的显存管理机制。

这一章，我们就来扒一扒PyTorch的显存管理内幕。你搞懂了这些，显存优化就成功了一半。

2.1 Tensor存储原理

Tensor在GPU上是怎么存的？这个问题我问过不少新人，答案五花八门。其实核心就两点：存储（Storage）和视图（View）。

2.1.1 存储与视图分离

每个Tensor背后都有一个Storage对象。Storage是真正的数据块，而Tensor只是对这个数据块的“视图”。

import torch

# 创建一个Tensor
x = torch.tensor([[1, 2, 3], [4, 5, 6]], device='cuda')
print(x.storage())  # 查看底层存储
# 输出: 1 2 3 4 5 6

# 取一个切片
y = x[0, :]  # 视图，不复制数据
print(y.storage() is x.storage())  # True，共享存储

我个人习惯用.storage()来检查是否真的复制了数据。很多OOM问题，就是因为不小心触发了数据复制。

核心要点：视图操作（切片、转置、expand等）不分配新显存。只有clone()、contiguous()等操作才会复制数据。

2.1.2 显存分配粒度

PyTorch分配显存不是按字节来的。它有一个最小分配单元，通常是512字节或2MB。你想想看，如果你只存一个float32（4字节），PyTorch也会给你分配512字节。

数据类型	元素大小	最小分配单元
float32	4字节	512字节
float16	2字节	512字节
int8	1字节	512字节

我在项目中遇到过一个问题：训练一个很小的模型，显存却占了1GB。排查了半天，发现是频繁创建小Tensor导致的碎片化。

2.2 CUDA缓存分配器

PyTorch的显存管理，底层用的是CUDA缓存分配器（CUDA Caching Allocator）。这个机制很有意思——它不直接向GPU申请显存，而是搞了个“缓存池”。

2.2.1 工作原理

说白了就是：PyTorch先向GPU申请一大块显存（比如2GB），然后自己管理这块显存。当你需要显存时，它从缓存池里切一块给你。当你释放Tensor时，它不还给GPU，而是放回缓存池。

import torch

# 查看缓存分配器状态
print(torch.cuda.memory_stats())

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

# 查看当前已分配显存
print(torch.cuda.memory_allocated())
# 查看缓存中保留的显存
print(torch.cuda.memory_reserved())

个人经验：我建议在训练循环开始前调用一次torch.cuda.empty_cache()，可以避免上一轮训练留下的缓存碎片。

2.2.2 缓存策略

缓存分配器有几个关键策略：

按需分配：需要多少显存，就从缓存池里切多少
延迟回收：释放的显存不立即还给GPU
碎片整理：当缓存池碎片太多时，会尝试合并相邻的空闲块

嗯，这里要注意：缓存分配器不是万能的。如果你频繁创建和销毁不同大小的Tensor，缓存池会变得支离破碎。

2.3 显存回收策略

显存回收，就是PyTorch怎么把不用的显存“还回来”。这里有两个机制：引用计数和垃圾回收。

2.3.1 引用计数

每个Tensor都有一个引用计数。当引用计数归零时，PyTorch会立即回收这块显存（放回缓存池）。

import torch

def create_tensor():
    x = torch.randn(1000, 1000, device='cuda')
    return x

# 函数返回后，x的引用计数归零
# 显存被回收（放回缓存池）
y = create_tensor()
print(torch.cuda.memory_allocated())  # 显存已被回收

2.3.2 垃圾回收

Python的垃圾回收器（GC）也会参与显存管理。当循环引用发生时，GC会帮忙清理。

我曾经踩过的坑：在训练循环中，如果某个Tensor被列表或字典引用着，即使你不再需要它，显存也不会被回收。我曾经因为这个原因，一个8GB的模型跑着跑着就OOM了。

2.4 torch.cuda最佳实践

说了这么多理论，来点实战的。以下是我总结的几条最佳实践：

2.4.1 监控显存

我习惯在训练脚本里加上显存监控：

import torch
import gc

def print_memory_usage():
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"已分配: {allocated:.2f} GB, 已预留: {reserved:.2f} GB")

# 在每个epoch结束时调用
for epoch in range(num_epochs):
    train_one_epoch()
    print_memory_usage()
    gc.collect()  # 手动触发垃圾回收
    torch.cuda.empty_cache()  # 清空缓存

2.4.2 避免显存碎片

显存碎片是OOM的隐形杀手。我建议：

预分配大Tensor：一次性分配好，不要频繁创建销毁
使用固定大小：batch size、序列长度尽量固定
复用Tensor：用zero_()或fill_()重置，而不是重新创建

# 不好的做法
for i in range(1000):
    x = torch.randn(1024, 1024, device='cuda')
    # 使用x...
    del x  # 频繁创建销毁

# 好的做法
x = torch.randn(1024, 1024, device='cuda')
for i in range(1000):
    x.normal_()  # 复用Tensor，只修改数据
    # 使用x...

2.4.3 显存分析工具

PyTorch自带了一些好用的工具：

工具	用途
torch.cuda.memory_summary()	打印详细的显存使用报告
torch.cuda.memory_snapshot()	获取显存快照，用于分析
torch.cuda.set_per_process_memory_fraction()	限制进程可用显存比例

我的习惯：在调试阶段，我会在关键位置插入torch.cuda.memory_summary()，看看显存到底被谁吃了。这个习惯帮我解决过不少OOM问题。

2.5 知识体系总览

下面这张图，是我自己整理的PyTorch显存管理知识体系。你把它记在脑子里，以后遇到显存问题就知道从哪下手了。

这张图把四个核心模块串起来了。你从Tensor存储原理入手，理解数据怎么存；再看缓存分配器，知道显存怎么管；然后学回收策略，明白显存怎么还；最后用最佳实践，把知识落地。

嗯，这一章的内容就到这里。记住：显存管理不是玄学，是有章可循的。你掌握了这些机制，以后遇到OOM就不会慌了。

📚 课程目录

第2章：PyTorch显存管理机制

2.1 Tensor存储原理

2.1.1 存储与视图分离

2.1.2 显存分配粒度

2.2 CUDA缓存分配器

2.2.1 工作原理

2.2.2 缓存策略

2.3 显存回收策略

2.3.1 引用计数

2.3.2 垃圾回收

2.4 torch.cuda最佳实践

2.4.1 监控显存

2.4.2 避免显存碎片

2.4.3 显存分析工具

2.5 知识体系总览