【Debug日志 | 模型loss不降调试】

import torch, torch.nn as nn, torch.nn.functional as F

class Net(nn.Module):
    def __init__(self): 
        super().__init__()
        self.feat = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
        )
        self.fc = nn.Linear(16, 1)

    def forward(self, x):
        f = self.feat(x)              # [B, 16]
        f_log = f.detach()            # ❌ 为了可视化，提前 detach
        # …日志里用到了 f_log …

        # 原地归一化（in-place）
        f /= (f.norm(dim=1, keepdim=True) + 1e-6)   # ❌ in-place 可能破坏版本计数
        return self.fc(f).squeeze(-1)

net = Net().cuda()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
scaler = torch.cuda.amp.GradScaler()

def l2_sp_regularizer(m: nn.Linear, tau=1e-4):
    # 为了稀疏化，做了个软阈值
    with torch.no_grad():
        m.weight.data = torch.clamp(m.weight.data, -1.0, 1.0)  # ❌ .data + 原地
    return tau * (m.weight.abs().mean())  # 看似没问题

for step in range(200):
    x = torch.randn(64, 3, 224, 224, device="cuda")
    y = (torch.rand(64, device="cuda") > 0.5).float()

    opt.zero_grad(set_to_none=True)
    with torch.cuda.amp.autocast(True):
        logits = net(x)
        base = F.binary_cross_entropy_with_logits(logits, y)
        reg  = l2_sp_regularizer(net.fc, 1e-4)
        loss = base + reg

    scaler.scale(loss).backward()   # ⬅️ 梯度经常是 None 或极小
    scaler.step(opt); scaler.update()

def tap_grad(t, name):
    t.retain_grad()
    def _hook(grad): print(f"[{name}] grad_norm={grad.norm().item():.4e}")
    t.register_hook(_hook)
    return t

with torch.cuda.amp.autocast(True):
    f = net.feat(x)
    tap_grad(f, "feat")     # ✅ 非叶子张量需要 retain_grad 才能看到 .grad
    logits = net.fc(f)
    loss = F.binary_cross_entropy_with_logits(logits, y)

loss.backward()
# 观察是否有打印；若无，则在更前面打点，直到发现哪一段“消失”

torch.autograd.set_detect_anomaly(True)

# ✅ 用 clone().detach() 生成只用于日志的副本，不参与计算
f = self.feat(x)                  # 参与反传
f_for_log = f.detach().clone()    # 仅用于可视化，别再写回去
# …用 f_for_log 画图/记录…

# ❌ f /= norm
# ✅
norm = (f.norm(dim=1, keepdim=True) + 1e-6)
f = f / norm

# ✅ 显式正则，进入计算图，由优化器“看得见”
def l2_sp_regularizer(m, tau=1e-4):
    return tau * m.weight.abs().mean()

# ✅ 如果要“硬裁剪”，用 optimizer hook 或 step 之后统一 clamp
@torch.no_grad()
def clamp_weights_(m: nn.Module, lo=-1.0, hi=1.0):
    for p in m.parameters():
        p.clamp_(lo, hi)

# 训练循环中：
scaler.scale(loss).backward()
scaler.step(opt); scaler.update()
clamp_weights_(net.fc)     # ✅ 在优化器更新后、no_grad 下原地裁剪