PyTorch 训练显存越跑越涨：隐式保留计算图导致 OOM

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

x = torch.randn(2000, 3, 224, 224)
y = (torch.rand(2000, 10) > 0.5).float()
loader = DataLoader(TensorDataset(x, y), batch_size=32, shuffle=True)

net = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(),
                    nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(16, 10)).cuda()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)

logits_buf, labels_buf, loss_hist = [], [], []  # ❌ 跨 step 的“粘性”列表

for step, (bx, by) in enumerate(loader, 1):
    bx, by = bx.cuda(), by.cuda()
    opt.zero_grad(set_to_none=True)
    logits = net(bx)                         # [B, 10], requires_grad=True
    loss = F.binary_cross_entropy_with_logits(logits, by)

    loss.backward()
    opt.step()

    # ❌ 直接把带梯度的张量放进列表，持有整条计算图
    logits_buf.append(logits)                # ← 泄漏点 1
    labels_buf.append(by)                    # ← 泄漏点 2
    loss_hist.append(loss)                   # ← 泄漏点 3：loss 张量也持图

    if step % 50 == 0:
        alloc = torch.cuda.memory_allocated() / 1024**2
        print(f"step {step}, mem={alloc:.1f} MB")

# 若后面还想做 epoch F1/PR 曲线，这些列表会继续增长并持有图，直至 OOM

def gpu_mb():
    return torch.cuda.memory_allocated() / 1024**2

# 在训练 loop 打点：
print(f"[dbg] before step={step}, mem={gpu_mb():.1f}MB")
...
print(f"[dbg] after  step={step}, mem={gpu_mb():.1f}MB")

import torch, torch.nn as nn, torch.nn.functional as F
logits_buf, labels_buf, loss_hist = [], [], []

for step, (bx, by) in enumerate(loader, 1):
    bx, by = bx.cuda(non_blocking=True), by.cuda(non_blocking=True)
    opt.zero_grad(set_to_none=True)

    with torch.cuda.amp.autocast(False):   # 可选：若用 AMP，保持默认策略即可
        logits = net(bx)
        loss = F.binary_cross_entropy_with_logits(logits, by)

    loss.backward()
    # 不要随意 retain_graph=True；若确需多次 backward，请定位到子图而非整图
    torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
    opt.step()

    # ✅ 仅保存“与训练解耦”的副本
    logits_buf.append(logits.detach().cpu())    # 切断计算图，移到 CPU
    labels_buf.append(by.detach().cpu())
    loss_hist.append(loss.item())               # 标量化

    if step % 50 == 0:
        print(f"step {step}, mem={torch.cuda.memory_allocated()/1024**2:.1f} MB")

# ✅ 指标计算放到 epoch 尾，并尽快释放 GPU 中间态
import torchmetrics
pred = torch.sigmoid(torch.cat(logits_buf)) > 0.5
tgt  = torch.cat(labels_buf).bool()
# ... 计算 F1/PR 等 ...
logits_buf.clear(); labels_buf.clear()          # 释放 CPU 内存引用
torch.cuda.empty_cache()                        # 可选：释放可缓存块（碎片化时有用）

# ✅ 验证阶段务必 no_grad
net.eval()
with torch.inference_mode():
    for bx, by in val_loader:
        # 验证不会增长显存
        _ = net(bx.cuda())
net.train()