4. 动态显存管理：torch.no_grad、del与gc.collect、in-place操作、显存碎片整理

显存这东西，说白了就是「僧多粥少」。模型越来越大，batch size 稍微调大一点就 OOM。很多同学跑来问我：「明明代码逻辑没问题，怎么跑着跑着显存就炸了？」

嗯，这里有个关键点——你释放显存的速度，赶不上你分配的速度。动态显存管理，就是教你怎么跟 GPU 这位「房东」搞好关系，该退房就退房，别占着茅坑不拉屎。

核心思想：显存不是无限的，但你可以通过「及时释放」、「原地修改」、「碎片整理」三个手段，把每一兆显存都用在刀刃上。

4.1 torch.no_grad：推理时别让梯度偷走你的显存

我个人习惯，只要做推理（inference），torch.no_grad() 就是肌肉记忆。为什么？因为 PyTorch 默认会为每个张量操作构建计算图，保留梯度信息。推理时你根本不需要反向传播，这些计算图就是纯浪费。

# 错误示范：推理时还开着梯度
def predict(model, x):
    return model(x)  # 这里会构建计算图，显存白白多占 30%-50%

# 正确做法
def predict(model, x):
    with torch.no_grad():
        return model(x)  # 不追踪梯度，显存占用直接砍半

我在项目中遇到过一位同事，跑 BERT 推理时 batch size 只能设 16，加了 torch.no_grad() 后直接飙到 32。你想想看，同样的硬件，性能翻倍，就少写了一行代码。

小技巧：如果你用 @torch.no_grad() 装饰器装饰整个函数，连缩进都省了。我一般推理脚本里直接写在函数头上，干净利落。

4.2 del 与 gc.collect：手动给 GPU 做「大扫除」

很多同学以为 Python 的垃圾回收会自动释放显存。天真了。Python 的引用计数机制对 CPU 内存有效，但 GPU 显存是「租来的地」，你不主动退租，房东（CUDA）不会帮你清。

什么时候需要手动 del？

循环中产生了中间张量，比如每轮迭代都创建了新的 feature map
大张量用完即弃，比如 loss 计算完就不再需要 logits
显存已经报警，但程序还没崩——这时候赶紧抢救

# 典型场景：循环中累积中间结果
intermediate_results = []
for i in range(100):
    feat = model.get_feature(x)  # 每个 feat 都占显存
    intermediate_results.append(feat)

# 处理完后，立刻释放
del intermediate_results
torch.cuda.empty_cache()  # 清空缓存，但不是必须的
gc.collect()              # 强制 Python 垃圾回收

我曾经在训练 GAN 时，生成器和判别器交替更新，中间特征图越积越多。一开始没加 del，跑到第 200 个 epoch 直接 OOM。后来在每个 batch 结束时手动 del fake_imgs, real_imgs, d_loss, g_loss，显存峰值降了 40%。

注意：torch.cuda.empty_cache() 不是银弹。它只是清空 CUDA 的缓存池，但不会释放已经被张量占用的显存。真正释放要靠 del 把引用计数归零。我见过有人疯狂调 empty_cache 但忘了 del，结果显存纹丝不动——白忙活。

4.3 in-place 操作：原地修改，省一半显存

in-place 操作，说白了就是「不另找房子，直接在原地址上改」。PyTorch 里带下划线的操作都是 in-place 的，比如 tensor.add_(x)、tensor.mul_(y)。

为什么省显存？因为普通操作会创建一个新张量，in-place 操作直接修改原张量。对于大张量，这能省下整整一份拷贝的显存。

# 普通操作：创建新张量，显存翻倍
a = torch.randn(10000, 10000).cuda()
b = a + 1  # 这里 a 和 b 同时存在，显存占用 2x

# in-place 操作：原地修改，显存不变
a = torch.randn(10000, 10000).cuda()
a.add_(1)  # 直接修改 a，不创建新张量

但这里有个大坑——in-place 操作会破坏计算图。如果你在训练时对需要梯度的张量用了 in-place，反向传播会直接报错。我刚开始用 PyTorch 时就踩过这个坑，debug 了一下午才发现是 relu_() 惹的祸。

经验法则：推理时放心用 in-place；训练时只在不需要梯度的张量上用（比如 with torch.no_grad() 包裹的代码块里）。

4.4 显存碎片整理：别让你的显存变成「蜂窝煤」

显存碎片，是个容易被忽视的问题。你想想看，频繁地分配和释放不同大小的张量，显存会被切成一块一块的「碎片」。就像硬盘碎片一样，虽然总容量够，但找不到一块连续的大空间来放你的大张量。

什么时候碎片化严重？

动态 batch size 或变长序列（比如 NLP 里的 padding）
频繁创建和销毁不同 shape 的中间张量
多模型交替推理/训练

我遇到过最夸张的一次，显存显示还剩 4GB，但分配一个 3GB 的张量竟然 OOM 了。查了半天，就是碎片太多，没有连续空间。

怎么整理？

预分配大块内存：一开始就申请好最大可能需要的显存，避免频繁小量分配
统一张量大小：尽量让所有 batch 的 shape 一致，减少碎片
使用 torch.cuda.empty_cache()：虽然不能完全解决碎片，但能整理 CUDA 的缓存池
重启进程：最暴力的方法，但有时候最有效。如果碎片化严重，直接 kill 进程重来

# 预分配策略：先创建一个最大尺寸的缓存张量
max_batch_size = 32
max_seq_len = 512
cache = torch.empty(max_batch_size, max_seq_len, 768).cuda()

# 后续使用时，直接切片操作，避免反复分配
for batch in dataloader:
    x = batch['input']  # shape 可能变化
    # 用 cache 的前 x.shape[0] 行，避免新分配
    cache[:x.shape[0], :x.shape[1], :] = x
    output = model(cache[:x.shape[0], :x.shape[1], :])

我的习惯：每次训练前先跑一个 dummy batch，把模型预热一下，顺便让 CUDA 把该分配的缓存都分配好。这样后续训练时显存分配更稳定，碎片也少。

知识体系总览

下面这张图，把本章的四个核心知识点串起来了。你可以把它当作显存优化的「作战地图」。

这四个手段，说白了就是「开源节流」——torch.no_grad 和 in-place 操作是节流，del 和碎片整理是开源。两手都要抓，两手都要硬。

最后说一句，显存优化不是一锤子买卖。你每次改模型结构、调 batch size，都可能打破之前的平衡。养成好习惯，把 del 和 no_grad 写进你的代码模板里，比事后救火强一百倍。

📚 课程目录

4. 动态显存管理：torch.no_grad、del与gc.collect、in-place操作、显存碎片整理

4.1 torch.no_grad：推理时别让梯度偷走你的显存

4.2 del 与 gc.collect：手动给 GPU 做「大扫除」

4.3 in-place 操作：原地修改，省一半显存

4.4 显存碎片整理：别让你的显存变成「蜂窝煤」

知识体系总览