4、张量生命周期管理：引用计数、延迟回收、梯度检查点、显存复用

显存碎片，说白了就是内存东一块西一块的空洞。

你分配一个张量，释放一个张量，再分配一个更大的——系统找不到连续空间，就报 OOM。我见过太多同学把 OOM 归咎于「显存不够」，其实很多时候是碎片在作祟。

这一章，我们聊聊怎么管好张量的生命周期。

4.1 引用计数：谁在用这个张量？

Python 的自动内存管理，靠的是引用计数。每个张量对象内部有个计数器，记录有多少变量在引用它。计数归零，内存就释放。

听起来很完美？其实坑不少。

⚠️ 循环引用陷阱

两个张量互相引用，计数永远不为零。Python 的垃圾回收器虽然能处理循环引用，但触发时机不确定。我曾在一次训练中，因为循环引用导致显存泄漏，跑了两天模型直接崩了。

我的习惯是：显式删除不再需要的张量。用 del 关键字，配合 torch.cuda.empty_cache() 及时回收。

import torch

# 显式删除大张量
big_tensor = torch.randn(10000, 10000).cuda()
# ... 使用完毕
del big_tensor
torch.cuda.empty_cache()  # 建议只在调试时用，生产环境慎用

💡 小技巧

用 sys.getrefcount(obj) 查看引用计数。我调试内存泄漏时，经常用这招定位谁没释放。

4.2 延迟回收：别急着扔

有些张量，你刚释放，下一秒又要用。频繁分配释放，碎片就来了。

延迟回收的思路是：把释放动作推迟到真正需要内存的时候。比如用一个缓存池，暂时保留已释放的张量空间。

我曾在项目中实现过一个简单的延迟回收器：

class DelayedReclaimer:
    def __init__(self, max_cache_size=10):
        self.cache = []
        self.max_cache_size = max_cache_size
    
    def release(self, tensor):
        if len(self.cache) < self.max_cache_size:
            self.cache.append(tensor)
        else:
            # 真的释放
            del tensor
    
    def reclaim(self):
        # 在需要显存时，强制回收
        for t in self.cache:
            del t
        self.cache.clear()
        torch.cuda.empty_cache()

嗯，这里要注意：延迟回收不能太久。否则显存占用居高不下，反而影响训练速度。

4.3 梯度检查点：用时间换空间

梯度检查点（Gradient Checkpointing）是显存优化的经典手段。说白了就是：不保存中间激活值，反向传播时重新计算。

为什么会这样？

前向传播时，每个层的输出（激活值）都要保存，供反向传播求梯度用。网络越深，保存的激活值越多。梯度检查点只保存少数关键节点的激活值，其他节点在反向传播时重新计算。

🔑 核心公式

显存节省 ≈ (1 - 1/N) × 原始显存，其中 N 是检查点间隔层数。

计算开销增加 ≈ 1 次额外前向传播。

import torch
from torch.utils.checkpoint import checkpoint

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(1024, 1024)
        self.layer2 = torch.nn.Linear(1024, 1024)
        self.layer3 = torch.nn.Linear(1024, 1024)
    
    def forward(self, x):
        # 对 layer2 使用梯度检查点
        x = self.layer1(x)
        x = checkpoint(self.layer2, x)  # 不保存中间激活
        x = self.layer3(x)
        return x

⚠️ 性能权衡

梯度检查点会增加约 20%-30% 的计算时间。我建议只在显存瓶颈时使用，并且只对最耗显存的层启用。

4.4 显存复用：一个萝卜一个坑

显存复用，就是让不同张量共享同一块物理显存。你想想看，如果两个张量不会同时存在，为什么不能共用空间？

PyTorch 的缓存分配器已经做了这件事。它维护了一个缓存池，释放的张量空间不会立即归还给系统，而是留在池子里，供后续分配使用。

我遇到过最典型的场景是：训练循环中，每个 batch 的 loss 张量。每次迭代都创建一个 loss，反向传播后释放。如果不复用，显存会持续增长。

# 错误做法：每次创建新张量
for batch in dataloader:
    loss = compute_loss(batch)  # 每次分配新显存
    loss.backward()
    # loss 释放，但显存碎片增加

# 正确做法：复用张量
loss_tensor = torch.zeros(1).cuda()
for batch in dataloader:
    loss_tensor.fill_(0)  # 复用同一块显存
    loss = compute_loss(batch, loss_tensor)
    loss.backward()

💡 我的经验

用 torch.cuda.memory_summary() 查看显存分配情况。我曾经靠这个命令，发现一个循环里反复创建了 1000 多个小张量，每个只占几 KB，但碎片率高达 40%。

4.5 综合策略：四管齐下

实际项目中，我通常把这四种方法组合使用：

策略	适用场景	代价
引用计数	所有场景，基础保障	无
延迟回收	频繁分配释放的小张量	少量显存占用
梯度检查点	深层网络，显存不足	20-30% 计算开销
显存复用	固定形状的中间张量	代码复杂度增加

我曾经在一个 72 层 Transformer 的训练中，通过组合使用梯度检查点和显存复用，把显存占用从 32GB 降到了 18GB，训练时间只增加了 15%。

嗯，这就是张量生命周期管理的核心。说白了，就是让每个字节的显存都物尽其用，不浪费、不闲置、不碎片化。

想要更多资料，联系可可