3、显存分配机制：PyTorch/TensorFlow的显存分配策略，CUDA context的显存开销

显存分配，说白了就是你的模型和数据怎么在GPU上安家。很多同学跑模型时遇到OOM（显存溢出），第一反应是“模型太大了”，其实很多时候是分配策略没搞对。我刚开始做训练时也踩过这个坑，明明算出来显存够用，一跑就崩，后来才发现是框架的分配机制在“捣鬼”。

3.1 CUDA Context：看不见的显存“钉子户”

你调用任何CUDA操作之前，系统会先创建一个CUDA context。这个东西是GPU的“运行环境”，包括内存管理、流调度、模块加载等。我个人的经验是，很多人忽略了context本身的显存开销。

为什么会这样？因为CUDA driver需要预留一部分显存来管理分配、释放、同步这些操作。你想想看，每次分配显存都要查表、更新元数据，这些数据总得有个地方放。

3.2 PyTorch 的显存分配策略

PyTorch用的是缓存分配器（Caching Allocator）。说白了就是：它不会每次需要显存都去问CUDA要，而是先自己囤一批，用完了也不急着还，留着下次用。

我个人习惯用 torch.cuda.memory_summary() 来查看分配情况。这个命令能告诉你当前缓存了多少、用了多少、碎片有多少。

import torch

# 查看当前显存分配情况
print(torch.cuda.memory_summary())

# 清空缓存（但不释放给系统）
torch.cuda.empty_cache()

# 查看当前已分配显存（单位：字节）
allocated = torch.cuda.memory_allocated()
cached = torch.cuda.memory_reserved()
print(f"已分配: {allocated / 1024**2:.2f} MB")
print(f"已缓存: {cached / 1024**2:.2f} MB")

这里有个关键点：empty_cache() 只是把缓存还给PyTorch自己的池子，并不是还给操作系统。所以你看 nvidia-smi 显示的显存占用可能没变。嗯，这里要注意，别被这个数字骗了。

PyTorch 分配策略的三大特点

特点	说明	影响
预分配大块	首次分配时直接要一大块（默认2MB起步）	减少CUDA调用次数，但可能浪费显存
流式分配	不同CUDA流使用独立的缓存池	多流并行时避免锁竞争
延迟释放	显存释放后放入缓存池，不立即归还	加速后续分配，但碎片化加剧

3.3 TensorFlow 的显存分配策略

TensorFlow 默认是“贪心模式”——它一上来就把所有可见的显存都占了。我第一次用TF时，看到 nvidia-smi 显示显存占用99%，以为程序出bug了，后来才知道这是它的“正常操作”。

TF 用的是 BFC（Best-Fit with Coalescing）分配器。名字听着复杂，其实就是“找最合适的空位放数据，放不下就合并碎片”。

import tensorflow as tf

# 方式1：限制显存增长（推荐）
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# 方式2：手动设置显存上限
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]  # 限制4GB
)

我个人建议用 set_memory_growth，它让TF按需分配，而不是一口气全占。你想想看，如果你机器上同时跑两个模型，一个占满显存，另一个就只能干瞪眼。

TF 和 PyTorch 的分配差异

对比项	PyTorch	TensorFlow
默认行为	按需分配，缓存复用	全量分配，贪心模式
碎片处理	依赖缓存池，碎片较多	BFC合并，碎片较少
多进程支持	每个进程独立分配器	支持进程间共享
显存回收	延迟回收，需手动empty_cache	自动回收，但较慢

3.4 显存分配的核心流程图

下面这张图是我自己总结的显存分配流程，从你调用 .cuda() 到数据真正落到显存里，中间经历了哪些步骤。

3.5 实战：如何监控和优化显存分配

我平时排查显存问题时，会按这个顺序来：

1. 先看context数量：用 nvidia-smi 看每个进程的显存占用，如果多个进程各占几百MB，大概率是context重复创建。
2. 再看缓存占用：用 torch.cuda.memory_reserved() 减去 memory_allocated()，差值就是缓存。如果缓存占比超过30%，说明碎片严重。
3. 最后看碎片率：PyTorch 1.9+ 提供了 torch.cuda.memory_stats()，里面有 num_alloc_retries 字段，这个值越大，碎片越严重。

嗯，显存分配这块其实还有很多细节，比如多流分配、异步释放、内存池大小调整等。但掌握了上面这些核心策略，大部分OOM问题你都能找到根因。记住一句话：显存不够，不一定是模型太大，也可能是分配策略没调好。