3、Tensor生命周期分析：张量的创建与销毁、引用计数与垃圾回收、Python GC对显存的影响

聊到显存复用，绕不开一个核心问题：张量什么时候生，什么时候死？

我刚开始做模型训练时，经常遇到显存莫名其妙涨上去，又莫名其妙降下来。后来才明白，这背后全是Tensor生命周期在作祟。说白了，你创建的每个张量，都在跟显存打交道。搞懂它的生死，你才能精准控制显存。

3.1 张量的创建：显存是怎么被占用的

每次你调用 torch.tensor() 或 torch.zeros()，PyTorch 都会在GPU上分配一块显存。这块显存有多大？取决于张量的形状和数据类型。

import torch

# 创建一个 float32 的张量，占用 4 * 1024 * 1024 = 4MB
x = torch.zeros(1024, 1024, dtype=torch.float32, device='cuda')
print(x.element_size() * x.numel())  # 输出 4194304 字节 ≈ 4MB

嗯，这里要注意：创建张量不等于立即分配显存。PyTorch 的缓存分配器会先看看有没有空闲块。如果有，直接复用；没有，才向GPU申请新内存。我在项目中遇到过，频繁创建小张量会导致分配器碎片化，显存利用率反而下降。

3.2 张量的销毁：引用计数与垃圾回收

张量什么时候销毁？这取决于Python的引用计数机制。

每个Python对象都有一个引用计数。当引用计数降到0，对象就会被销毁，底层显存也会被释放。但注意，这里的“释放”是指归还给PyTorch的缓存池，而不是直接还给GPU。

import torch

def create_tensor():
    # 引用计数为1
    x = torch.randn(1000, 1000, device='cuda')
    # 函数结束，x的引用计数变为0，张量被销毁
    # 显存回到缓存池

create_tensor()
# 此时显存并没有真正释放，只是回到了缓存池

我见过不少同学以为函数结束显存就释放了，其实不然。PyTorch的缓存分配器会保留这些显存块，方便下次快速分配。如果你想强制释放，得调用 torch.cuda.empty_cache()。

3.3 Python GC 对显存的影响

Python的垃圾回收器（GC）主要处理循环引用。对于张量来说，如果两个张量互相引用，引用计数就永远不会归零。这时候GC就派上用场了。

但GC有个问题：它不会立即触发。GC默认在分配对象达到一定阈值时才运行。这意味着，即使张量已经没用了，显存也可能被占用很久。

import gc
import torch

# 模拟循环引用
a = torch.randn(1000, 1000, device='cuda')
b = torch.randn(1000, 1000, device='cuda')
a.ref = b
b.ref = a

# 删除引用
del a
del b

# 此时显存还没释放，因为循环引用
# 手动触发GC
gc.collect()
# 显存才真正回到缓存池

你想想看，如果训练脚本里不小心产生了循环引用，GC又不及时触发，显存就会一直涨。我建议在关键位置手动调用 gc.collect()，尤其是处理大张量时。

3.4 生命周期可视化：一张图看懂

为了让你更直观地理解，我画了一张流程图。它展示了张量从创建到销毁的完整路径，以及显存分配和回收的时机。

从这张图能看出：张量销毁后，显存并不一定立即释放。它可能回到缓存池，也可能被GC延迟回收。这就是为什么你有时候看到显存占用居高不下。

3.5 实战建议：如何管理Tensor生命周期

基于上面的分析，我总结了几条实战建议：

• 及时删除无用张量：用 del 显式删除，让引用计数尽快归零。
• 避免循环引用：不要在张量上挂载互相引用的属性。如果必须用，记得手动 gc.collect()。
• 使用上下文管理器：比如 with torch.no_grad()，可以避免中间张量被保留。
• 监控引用计数：用 sys.getrefcount(x) 检查张量的引用情况，排查泄漏。

好了，这一章就到这里。记住，显存管理不是玄学，是科学。搞懂生命周期，你就成功了一半。