【Debug日志 | 断点训练异常】

import argparse, math, os, torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)

def make_loader(n=4096, bs=64):
    X = torch.randn(n, 10)
    y = (X[:, 0] + 0.6 * X[:, 1] > 0).long()
    ds = torch.utils.data.TensorDataset(X, y)
    return torch.utils.data.DataLoader(ds, batch_size=bs, shuffle=True, drop_last=True)

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.m = nn.Sequential(nn.Linear(10,64), nn.ReLU(), nn.Linear(64,2))
    def forward(self, x): return self.m(x)

def build_scheduler(optimizer, total_steps, warmup_steps=50):
    def lr_lambda(step):
        if step < warmup_steps: return (step + 1) / warmup_steps
        t = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, t)))
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

def save_ckpt(path, model, optimizer, scheduler, scaler, global_step):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "scaler": scaler.state_dict() if scaler is not None else None,
        "global_step": global_step,
    }, path)

def load_ckpt(path, model, optimizer=None, scheduler=None, scaler=None):
    ckpt = torch.load(path, map_location="cpu")
    model.load_state_dict(ckpt["model"])
    if optimizer is not None and "optimizer" in ckpt: optimizer.load_state_dict(ckpt["optimizer"])
    if scheduler is not None and "scheduler" in ckpt: scheduler.load_state_dict(ckpt["scheduler"])
    if scaler is not None and ckpt.get("scaler") is not None: scaler.load_state_dict(ckpt["scaler"])
    return ckpt.get("global_step", 0)

def train(phase, bug, steps):
    device = "cpu"
    model = TinyNet().to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    scaler = None  # GPU+AMP 时可替换为 GradScaler()
    total_steps = 400
    scheduler = build_scheduler(optimizer, total_steps)

    loader = make_loader()
    it = iter(loader)
    global_step = 0
    ckpt_path = "ckpts/demo.pt"

    if phase == "resume":
        if bug:
            _ = load_ckpt(ckpt_path, model)   # 只加载模型，错误示范
            global_step = 0
        else:
            global_step = load_ckpt(ckpt_path, model, optimizer, scheduler, scaler)

    for _ in range(steps):
        try:   x, y = next(it)
        except StopIteration:
            it = iter(loader); x, y = next(it)
        logits = model(x)
        loss = F.cross_entropy(logits, y)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        scheduler.step()
        global_step += 1

        if (global_step % 25) == 0:
            lr = optimizer.param_groups[0]["lr"]
            acc = (logits.argmax(1) == y).float().mean().item()
            tag = f"[{phase}|{'BUG' if bug else 'OK'}]"
            print(f"{tag} step={global_step:04d} lr={lr:.4f} loss={loss.item():.3f} acc={acc:.3f}")

    if phase == "pretrain":
        save_ckpt(ckpt_path, model, optimizer, scheduler, scaler, global_step)

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--phase", choices=["pretrain","resume"], required=True)
    ap.add_argument("--bug", choices=["on","off"], default="on")
    ap.add_argument("--steps", type=int, default=200)
    args = ap.parse_args()
    train(args.phase, args.bug=="on", args.steps)

# 阶段1：从头训练并保存检查点
python resume_bug_demo.py --phase pretrain --steps 200

# 阶段2a：错误续训（只加载模型权重）
python resume_bug_demo.py --phase resume --bug on --steps 200

# 阶段2b：正确续训（加载模型+优化器+调度器+scaler+全局步）
python resume_bug_demo.py --phase resume --bug off --steps 200

print("optimizer_has_state:", any(len(s)>0 for s in optimizer.state.values()))
print("scheduler_last_epoch:", getattr(scheduler, "last_epoch", None))
print("global_step:", global_step)

state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict() if scheduler else None,
    "scaler": scaler.state_dict() if scaler else None,
    "global_step": global_step,
    "epoch": epoch,
}
torch.save(state, ckpt_path)

state = torch.load(ckpt_path, map_location=device)
model.load_state_dict(state["model"])
if optimizer and state.get("optimizer"): optimizer.load_state_dict(state["optimizer"])
if scheduler and state.get("scheduler"): scheduler.load_state_dict(state["scheduler"])
if scaler and state.get("scaler"): scaler.load_state_dict(state["scaler"])
global_step = state.get("global_step", 0)
epoch = state.get("epoch", 0)

if (global_step % log_every) == 0: ...
if (global_step % eval_every) == 0: ...
if (global_step % save_every) == 0: ...