PyTorch显存管理机制：缓存分配器、CUDA stream、tensor生命周期

大家好，我是你们的老朋友。今天咱们来聊聊PyTorch的显存管理机制。说实话，这玩意儿我刚开始接触的时候也一头雾水——明明代码写得好好的，怎么跑着跑着就OOM了？后来踩了不少坑，才慢慢摸清楚门道。

PyTorch的显存管理，说白了就是三件事：谁分配显存、怎么分配、什么时候释放。对应到具体技术上，就是缓存分配器、CUDA stream和tensor生命周期。咱们一个一个来拆解。

1. 缓存分配器：PyTorch的“内存管家”

先问大家一个问题：当你创建一个tensor并放到GPU上时，PyTorch真的会立刻向CUDA申请显存吗？

答案是：不一定。

PyTorch内部有一个缓存分配器（Caching Allocator），它就像一个精明的管家。你每次要显存，它先看看自己手里有没有现成的。如果有，直接给你；如果没有，它才去向CUDA申请。

我举个例子。假设你训练一个模型，每次迭代都要创建一个临时tensor。如果没有缓存分配器，每次都要跟CUDA打交道——申请、释放、再申请、再释放。这就像你每次喝水都要去井边打水，效率极低。

有了缓存分配器呢？它会把释放的显存先“囤”起来。下次你再要，直接从囤货里拿。速度快得多。

  核心机制：PyTorch的缓存分配器使用cudaMalloc和cudaFree来管理底层显存，但上层做了缓存池化。默认情况下，它会缓存所有已释放的显存块，直到进程结束。

这里有个坑。我曾经调试一个模型，发现显存占用只增不减。明明我手动删除了tensor，还调了torch.cuda.empty_cache()，但显存就是下不来。为什么？

因为empty_cache()只是清空缓存分配器里“空闲”的块，但已经分配给tensor的显存它不管。说白了，它释放的是“管家囤的货”，不是“你手里正在用的”。

小技巧：如果你发现显存占用异常，可以用torch.cuda.memory_summary()查看详细分配情况。这个函数会列出所有活跃的分配块，帮你定位问题。

2. CUDA Stream：并行执行的“流水线”

接下来聊聊CUDA stream。这个概念很多人容易忽略，但它对显存管理影响很大。

CUDA stream，你可以把它想象成一条执行流水线。默认情况下，PyTorch所有操作都在同一个stream（默认流）上执行。这意味着操作是串行的——一个做完，下一个才能开始。

但如果你有多个stream呢？它们可以并行执行。比如，你可以在stream A上做前向计算，同时在stream B上做数据加载。这样能提高GPU利用率。

不过，stream和显存管理有什么关系？

关系大了。每个stream都有自己的上下文。如果你在stream A上分配了一个tensor，然后在stream B上使用它，PyTorch会自动插入同步操作。这个同步会阻塞stream B，直到stream A完成。

我遇到过一个问题：用多个stream做数据预处理，结果显存暴涨。后来发现，是因为stream之间的同步导致tensor生命周期延长了——本该释放的tensor，因为被另一个stream引用着，迟迟释放不掉。

避坑指南：我曾经在训练循环里用多个stream做数据增强，结果显存从8G飙到12G。排查后发现，是stream同步导致临时tensor无法及时释放。解决方案是：在每次迭代结束时，显式调用torch.cuda.synchronize()，确保所有stream都完成，再清理缓存。

这里给个简单的代码示例，展示如何创建和使用stream：

import torch

# 创建两个stream
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

# 在stream上执行操作
with torch.cuda.stream(s1):
    # 这个tensor在s1上分配
    a = torch.randn(1000, 1000, device='cuda')
    b = a * 2  # 计算也在s1上

with torch.cuda.stream(s2):
    # 这个tensor在s2上分配
    c = torch.randn(1000, 1000, device='cuda')
    d = c + 1

# 默认stream（主流）等待所有stream完成
torch.cuda.synchronize()

嗯，这里要注意：不要滥用stream。如果你的模型不大，单stream就够用了。多stream带来的同步开销，有时候反而会拖慢速度。

3. Tensor生命周期：从生到死

最后，咱们聊聊tensor的生命周期。这是显存管理的核心——你只有理解了tensor什么时候生、什么时候死，才能精准控制显存。

一个tensor的生命周期分为三个阶段：

分配：当你创建tensor并放到GPU上时，PyTorch从缓存分配器里拿一块显存。
使用：tensor参与计算，显存被占用。
释放：当tensor的引用计数归零时，Python的垃圾回收器会回收它，显存回到缓存分配器。

这里有个关键点：引用计数。Python的垃圾回收机制是引用计数为主，标记清除为辅。只要还有变量引用着tensor，它就不会被释放。

我见过不少新手写这样的代码：

def train_step(x):
    # 每次调用都创建新tensor
    y = model(x)
    loss = y.sum()
    loss.backward()
    # y和loss在这里不会被释放，因为函数返回后还有引用？
    return loss.item()

实际上，y和loss在函数返回后，如果没有外部引用，就会被回收。但问题在于，loss.backward()会保留计算图，直到你调用optimizer.step()和optimizer.zero_grad()。如果计算图没释放，中间tensor就活着。

  关键结论：tensor的生命周期 = Python引用 + 计算图依赖。只要计算图还在，中间tensor就不会释放。所以，及时切断计算图（比如用.detach()）是显存优化的关键。

我自己的习惯是：在训练循环里，每次迭代结束后，显式调用del删除不再需要的tensor，再调torch.cuda.empty_cache()。虽然有点啰嗦，但能有效防止显存泄漏。

下面这张图展示了tensor从分配到释放的完整流程：

4. 实战建议：如何用好这三板斧

说了这么多，咱们总结一下实战中怎么用：

机制	常见问题	我的建议
缓存分配器	显存只增不减	用`memory_summary()`定位；调`empty_cache()`释放空闲块
CUDA Stream	多stream导致显存暴涨	控制stream数量；每次迭代结束同步一次
Tensor生命周期	计算图残留	及时`.detach()`；用`del`手动删除；调`gc.collect()`

我的个人习惯：每次训练迭代结束时，我会写这样一段代码：
del loss, outputs
torch.cuda.empty_cache()
import gc; gc.collect()
虽然有点“暴力”，但能保证显存不泄漏。特别是跑长训练时，这个习惯救了我好几次。

好了，今天的内容就到这里。PyTorch的显存管理，说白了就是理解这三个机制怎么配合。你只要记住：分配靠缓存、执行靠stream、释放靠引用，大部分显存问题都能迎刃而解。

想要更多资料，联系可可

📚 课程目录

PyTorch显存管理机制：缓存分配器、CUDA stream、tensor生命周期

1. 缓存分配器：PyTorch的“内存管家”

2. CUDA Stream：并行执行的“流水线”

3. Tensor生命周期：从生到死

4. 实战建议：如何用好这三板斧