"""真实类别增量学习：冻结方案后训练，保留全部种子、支路和错误。""" from argparse import ArgumentParser from copy import deepcopy import json from pathlib import Path import platform import random import time import numpy as np import sklearn from sklearn.datasets import load_digits import torch from torch import nn from torch.nn import functional as F import continual_learning as lab def data(): x, y = load_digits(return_X_y=True) raw = json.dumps({"pixels": x.astype(int).tolist(), "labels": y.tolist()}, separators=(",", ":")).encode() return torch.tensor(x/16, dtype=torch.float32), torch.tensor(y), lab.sha256(raw).hexdigest() def make_plan(): _, y, digest = data() rng = random.Random(20261020) groups = {f"{split}_{age}": [] for split in ["train", "dev", "test"] for age in ["old", "new"]} memory = [] for label in range(10): ids = [i for i, value in enumerate(y.tolist()) if value == label] rng.shuffle(ids) age = "old" if label < 5 else "new" groups[f"train_{age}"].extend(ids[:100]) groups[f"dev_{age}"].extend(ids[100:120]) groups[f"test_{age}"].extend(ids[120:150]) if label < 5: memory.extend(ids[:10]) return dict(version=1, date="2026-10-20", dataset_sha256=digest, ids={k: sorted(v) for k, v in groups.items()}, memory=sorted(memory), seeds=[11, 22, 33], arms=lab.ARMS, hidden=48, lr=.003, old_steps=240, new_steps=240, batch=32, replay_old_per_batch=8, dev_every=20, fisher="predictive expectation over 10 labels on 50 old inputs", ewc_lambda={"ewc100": 100., "ewc1000": 1000.}, selection="all arms, all seeds, last checkpoint; no dev/test parameter search", evaluation="one shared 10-class head, no task label or output mask") def model(plan): return nn.Sequential(nn.Linear(64, plan["hidden"]), nn.Tanh(), nn.Linear(plan["hidden"], 10)) def metrics(logits, labels): z, y = torch.tensor(logits), torch.tensor(labels) correct = z.argmax(-1) == y old = y < 5 return dict(n=len(y), correct=int(correct.sum()), accuracy=correct.float().mean().item(), old_n=int(old.sum()), old_correct=int(correct[old].sum()), old_accuracy=correct[old].float().mean().item(), new_n=int((~old).sum()), new_correct=int(correct[~old].sum()), new_accuracy=correct[~old].float().mean().item(), nll=F.cross_entropy(z, y).item()) @torch.no_grad() def evaluate(net, x, y, ids): logits, labels = net(x[ids]).tolist(), y[ids].tolist() return dict(ids=ids, labels=labels, logits=logits, metrics=metrics(logits, labels)) def fisher_diagonal(net, x): """对输入均匀平均，对标签按旧模型分布精确求和；不是批梯度的平方。""" params = list(net.parameters()) importance = [torch.zeros_like(p) for p in params] for example in x: logp = net(example.unsqueeze(0)).log_softmax(-1)[0] probabilities = logp.detach().exp() for label in range(len(probabilities)): gradients = torch.autograd.grad(logp[label], params, retain_graph=label+1 < len(probabilities)) for total, gradient in zip(importance, gradients): total.add_(probabilities[label]*gradient.square()/len(x)) return importance def penalty(net, anchor, importance): return sum((f*(p-old).square()).sum() for p, old, f in zip(net.parameters(), anchor, importance))/2 def train(net, x, y, pool, plan, seed, steps, dev_ids, replay=None, importance=None, strength=0.): generator = torch.Generator().manual_seed(seed+10000) old_generator = torch.Generator().manual_seed(seed+20000) anchor = [p.detach().clone() for p in net.parameters()] optimizer = torch.optim.Adam(net.parameters(), lr=plan["lr"]) # 每个阶段重新建立优化器。 history = [] curve = [dict(step=0, **evaluate(net, x, y, dev_ids)["metrics"])] for step in range(steps): # 先为各支路抽同样的 32 个新样本；回放支路只取前 24 个。 positions = torch.randint(len(pool), (plan["batch"],), generator=generator) ids = [pool[i] for i in positions.tolist()] if replay is not None: n_old = plan["replay_old_per_batch"] old_positions = torch.randint(len(replay), (n_old,), generator=old_generator) ids = ids[:-n_old]+[replay[i] for i in old_positions.tolist()] batch = torch.tensor(ids) ce = F.cross_entropy(net(x[batch]), y[batch]) regularizer = penalty(net, anchor, importance) if importance is not None else torch.tensor(0.) loss = ce+strength*regularizer optimizer.zero_grad() loss.backward() optimizer.step() history.append(dict(step=step+1, ids=ids, cross_entropy=ce.item(), penalty=regularizer.item(), loss=loss.item())) if (step+1) % plan["dev_every"] == 0: curve.append(dict(step=step+1, **evaluate(net, x, y, dev_ids)["metrics"])) return history, curve def run(plan_path, out): out.mkdir(parents=True, exist_ok=False) plan = json.loads(plan_path.read_text()) assert plan == make_plan() torch.set_num_threads(1) torch.use_deterministic_algorithms(True) x, y, raw_hash = data() assert raw_hash == plan["dataset_sha256"] split_ids = {s: sorted(plan["ids"][f"{s}_old"]+plan["ids"][f"{s}_new"]) for s in ["dev", "test"]} record = dict(complete=False, plan_sha256=lab.digest(plan_path), source_sha256={p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]}, environment=dict(python=platform.python_version(), torch=torch.__version__, numpy=np.__version__, sklearn=sklearn.__version__, platform=platform.platform(), threads=1, device="cpu"), parameters=sum(p.numel() for p in model(plan).parameters()), bases=[], runs=[]) started = time.monotonic() def checkpoint(): temporary = out/"record.pending.json" temporary.write_text(json.dumps(record, ensure_ascii=False, separators=(",", ":"))+"\n") temporary.replace(out/"record.json") def save_model(net, name, evaluations): path = out/f"{name}.pt" torch.save({"plan": plan, "state_dict": net.state_dict()}, path) restored = model(plan) restored.load_state_dict(torch.load(path, weights_only=True)["state_dict"]) assert all(evaluate(restored, x, y, split_ids[s]) == ev for s, ev in evaluations.items()) return dict(file=path.name, sha256=lab.digest(path), reload_equal=True) for seed in plan["seeds"]: torch.manual_seed(seed) old_model = model(plan) history, curve = train(old_model, x, y, plan["ids"]["train_old"], plan, seed, plan["old_steps"], split_ids["dev"]) evaluations = {s: evaluate(old_model, x, y, ids) for s, ids in split_ids.items()} fisher_started = time.monotonic() importance = fisher_diagonal(old_model, x[plan["memory"]]) fisher_seconds = time.monotonic()-fisher_started base = dict(seed=seed, history=history, curve=curve, evaluations=evaluations, fisher=[f.tolist() for f in importance], fisher_seconds=fisher_seconds, checkpoint=save_model(old_model, f"seed-{seed}-old", evaluations)) record["bases"].append(base) checkpoint() for arm in plan["arms"]: net = deepcopy(old_model) replay = plan["memory"] if arm == "replay50" else plan["ids"]["train_old"] if arm == "replay500" else None strength = plan["ewc_lambda"].get(arm, 0.) history, curve = train(net, x, y, plan["ids"]["train_new"], plan, seed+30000, plan["new_steps"], split_ids["dev"], replay, importance if strength else None, strength) evaluations = {s: evaluate(net, x, y, ids) for s, ids in split_ids.items()} record["runs"].append(dict(seed=seed, arm=arm, history=history, curve=curve, old_test_before=base["evaluations"]["test"]["metrics"]["old_accuracy"], evaluations=evaluations, checkpoint=save_model(net, f"seed-{seed}-{arm}", evaluations))) checkpoint() m = evaluations["test"]["metrics"] print(f"seed={seed} arm={arm} old={m['old_correct']}/{m['old_n']} " f"new={m['new_correct']}/{m['new_n']}", flush=True) record["wall_seconds"] = time.monotonic()-started record["complete"] = True checkpoint() print(f"完成 3 次旧类训练与 15 次后续训练：{out}") def main(): parser = ArgumentParser(description=__doc__) parser.add_argument("--freeze", type=Path) parser.add_argument("--run", action="store_true") parser.add_argument("--plan", type=Path, default=lab.PLAN) parser.add_argument("--out-dir", type=Path) args = parser.parse_args() if args.freeze: lab.save_new(args.freeze, make_plan()) print(f"已冻结划分、回放样本与全部训练配置：{args.freeze}") return if not args.run: print("默认不训练。使用 --run --out-dir 新目录，才运行全部 18 次训练。") return if args.out_dir is None: parser.error("需要尚不存在的 --out-dir") run(args.plan, args.out_dir) if __name__ == "__main__": main()