第四章：张量生命周期管理：Python对象引用与显存泄漏的真相

说实话，我见过太多同学在显存优化上栽跟头了。

很多人以为显存溢出就是模型太大、batch太大。其实不然。我遇到过最头疼的问题，反而是那些「看不见」的显存泄漏——Python对象引用把张量死死拽住，GPU想回收都回收不了。

今天咱们就来扒一扒，张量到底是怎么生、怎么死、又是怎么被「赖着不走」的。

4.1 张量的生与死：谁在掌控它的生命周期？

每个张量在GPU里都有两个阶段：分配和释放。

分配好理解，你调用 torch.tensor() 或者模型前向传播时，CUDA就会在显存里划一块地给它。

释放呢？Python的垃圾回收器（GC）说了算。但GC不是实时响应的，它有自己的节奏。这就埋下了隐患。

我举个例子你就明白了：

import torch

def leak_example():
    # 创建一个张量，占用1GB显存
    big_tensor = torch.randn(1024, 1024, 256, device='cuda')
    # 函数结束，big_tensor按理说应该被销毁
    # 但如果有人还在引用它...那就麻烦了
    return big_tensor  # 返回引用，外部接着用

# 调用函数，拿到引用
ref = leak_example()
# 此时ref还活着，显存不会释放

你看，只要 ref 这个变量还在，那个1GB的张量就赖在显存里不走。

4.2 最常见的显存泄漏陷阱

我个人习惯把泄漏场景分成三类。咱们一个一个说。

陷阱一：列表/字典里的「幽灵引用」

这是最隐蔽的。你想想看，如果你把一个张量塞进列表里，然后忘了清理列表——

class ModelWrapper:
    def __init__(self):
        self.cache = []  # 用来缓存中间结果
    
    def forward(self, x):
        # 每次前向都往cache里塞张量
        self.cache.append(x.detach().clone())
        return x * 2
    
    def clear_cache(self):
        self.cache.clear()  # 记得手动清理！

我在项目中遇到过这种情况：一个训练循环跑了1000步，cache里攒了1000个中间张量。显存从8GB一路飙到80GB，然后OOM了。

嗯，排查了整整一下午才发现是这里的问题。

陷阱二：闭包与lambda的「隐形手」

闭包会捕获外部变量。如果捕获的是张量，那这个张量就永远别想被释放了。

def create_callback(tensor):
    # lambda捕获了tensor的引用
    return lambda: tensor.sum().item()

big_tensor = torch.randn(10000, 10000, device='cuda')
cb = create_callback(big_tensor)
# 即使del big_tensor，cb里还抓着引用不放
del big_tensor
# 显存依然被占用！

我曾经调试过一个线上服务，显存一天涨10%。最后发现是回调函数列表里攒了几万个lambda，每个都抓着一个小张量不放。日积月累，就炸了。

陷阱三：DataLoader的worker进程

这个坑比较深。PyTorch的DataLoader如果设置 num_workers > 0，每个worker都会有自己的Python环境。如果worker里创建了CUDA张量并且没有正确清理——

# 错误示范：在worker里创建CUDA张量
def bad_collate_fn(batch):
    # 每个worker都会创建这个张量
    temp = torch.zeros(1000, 1000, device='cuda')
    # 函数结束，temp应该被释放
    # 但CUDA上下文还在，显存不会完全归还
    return batch

loader = DataLoader(dataset, batch_size=32, 
                    num_workers=4, 
                    collate_fn=bad_collate_fn)

说白了，worker里尽量不要碰CUDA。如果非要用，记得在worker退出前手动 torch.cuda.empty_cache()。

4.3 如何诊断显存泄漏？

光知道理论不够，得会动手查。我常用的三板斧：

方法	工具	适用场景
显存快照	torch.cuda.memory_summary()	快速查看当前显存分配情况
对象追踪	gc.get_objects()	找出哪些Python对象还活着
引用计数	sys.getrefcount()	确认某个张量被引用了多少次

我最常用的是 gc.get_objects()。配合类型过滤，能快速定位问题：

import gc
import torch

def find_tensor_leaks():
    # 获取所有存活对象
    objects = gc.get_objects()
    tensor_count = 0
    total_memory = 0
    
    for obj in objects:
        if torch.is_tensor(obj) and obj.is_cuda:
            tensor_count += 1
            total_memory += obj.numel() * obj.element_size()
    
    print(f"存活CUDA张量: {tensor_count} 个")
    print(f"总占用: {total_memory / 1024**3:.2f} GB")
    
    # 如果数量异常，打印前几个的引用链
    if tensor_count > 100:
        print("警告：张量数量过多，可能存在泄漏！")

4.4 最佳实践：让张量「死得其所」

说了这么多问题，咱们聊聊怎么解决。我总结了几条铁律：

1. 显式清理容器：列表、字典、队列用完就 .clear() 或 del
2. 避免闭包捕获：如果必须用回调，传张量索引而不是张量本身
3. 监控显存趋势：每100步打印一次 torch.cuda.memory_allocated()
4. 使用上下文管理器：自定义一个 TensorManager 来统一管理生命周期

class TensorManager:
    """一个简单的张量生命周期管理器"""
    def __init__(self):
        self._tensors = []
    
    def register(self, tensor):
        self._tensors.append(tensor)
        return tensor
    
    def release_all(self):
        for t in self._tensors:
            del t
        self._tensors.clear()
        torch.cuda.empty_cache()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.release_all()

# 使用示例
with TensorManager() as mgr:
    a = mgr.register(torch.randn(1000, 1000, device='cuda'))
    b = mgr.register(torch.randn(1000, 1000, device='cuda'))
    # 做点什么...
# 退出with块时，a和b自动释放

4.5 一张图看懂张量生命周期

下面这张图，是我自己画的知识结构。它把整个生命周期串起来了：

这张图把整个流程串起来了。你从左边开始看：创建张量 → 引用计数管理 → 释放显存。中间那条黄色的泄漏路径，就是咱们要重点防范的。

4.6 写在最后

张量生命周期管理，说白了就是管好「引用」这根线。

线断了，张量就死了，显存就释放了。线不断，哪怕你 del 了变量名，只要还有别的引用拽着，显存就还在。

我个人习惯在每次训练循环结束时，都跑一遍 gc.collect() 和 torch.cuda.empty_cache()。虽然不能根治泄漏，但至少能让显存「喘口气」。

记住一句话：你不管张量的死活，OOM就会来管你。

想要更多资料，联系可可