"""固定预算的真实 digits 筛选对照；--freeze 保存计划，--run 才训练。""" 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 from sklearn.model_selection import train_test_split import torch from torch import nn from torch.nn import functional as F import data_efficiency 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(): x, y, raw_hash = data() pool, test = train_test_split(np.arange(len(y)), test_size=300, stratify=y.numpy(), random_state=20261020) pool, dev = train_test_split(pool, test_size=200, stratify=y[pool].numpy(), random_state=20261021) rng = random.Random(20261022) selector_ids, candidates = [], [] for label in range(10): ids = [int(i) for i in pool if int(y[i]) == label] rng.shuffle(ids) selector_ids.extend(ids[:10]) count = 100 if label < 8 else 20 assert len(ids) >= 10+count candidates.extend(ids[10:10+count]) selector_ids.sort() candidates.sort() observed = {str(i): int(y[i]) for i in candidates} changed = sorted(rng.sample(candidates, round(.08*len(candidates)))) for i in changed: observed[str(i)] = (observed[str(i)]+rng.randrange(1, 10)) % 10 assert min(sum(v == k for v in observed.values()) for k in range(10)) >= 16 return {"version": 1, "date": "2026-10-20", "dataset_sha256": raw_hash, "selector_ids": selector_ids, "candidate_ids": candidates, "dev_ids": sorted(map(int, dev)), "test_ids": sorted(map(int, test)), "observed_labels": observed, "changed_ids": changed, "seeds": [11, 22, 33], "arms": lab.ARMS, "subset_size": 160, "hidden": 32, "lr": .003, "batch_size": 32, "steps": 240, "selector_seed": 20261023, "selector_steps": 200, "curriculum_limits": [[60, 40], [120, 80], [240, 160]], "selection": "last checkpoint, all arms and seeds; no dev/test selection", "corruption": "67 controlled label changes; originals used only for audit/evaluation"} def model(plan): return nn.Sequential(nn.Linear(64, plan["hidden"]), nn.Tanh(), nn.Linear(plan["hidden"], 10)) def train(net, x, labels, ids, plan, seed, steps, curriculum=False): generator = torch.Generator().manual_seed(seed+10000) optimizer = torch.optim.Adam(net.parameters(), lr=plan["lr"]) history = [] for step in range(steps): limit = len(ids) if curriculum: limit = next(n for until, n in plan["curriculum_limits"] if step < until) positions = torch.randint(limit, (plan["batch_size"],), generator=generator) batch = torch.tensor([ids[i] for i in positions.tolist()]) loss = F.cross_entropy(net(x[batch]), labels[batch]) optimizer.zero_grad() loss.backward() optimizer.step() history.append({"step": step+1, "pool_size": limit, "ids": batch.tolist(), "loss": loss.item()}) return history def choose(plan, observed, scores, seed, arm): candidates = plan["candidate_ids"] if arm == "all": return candidates[:] rng = random.Random(seed+20000) if arm in {"random", "curriculum"}: chosen = rng.sample(candidates, plan["subset_size"]) return sorted(chosen, key=lambda i: (scores[i], i)) if arm == "curriculum" else sorted(chosen) if arm in {"easy", "hard"}: direction = 1 if arm == "easy" else -1 return sorted(candidates, key=lambda i: (direction*scores[i], i))[:plan["subset_size"]] if arm == "balanced": result = [] for label in range(10): group = [i for i in candidates if int(observed[i]) == label] result.extend(sorted(group, key=lambda i: (scores[i], i))[:16]) assert len(result) == plan["subset_size"] return sorted(result) raise ValueError(arm) def metrics(logits, labels): z, y = torch.tensor(logits), torch.tensor(labels) correct = z.argmax(-1) == y tail = (y == 8) | (y == 9) return {"n": len(y), "correct": int(correct.sum()), "accuracy": correct.float().mean().item(), "tail_n": int(tail.sum()), "tail_correct": int(correct[tail].sum()), "tail_accuracy": correct[tail].float().mean().item(), "nll": F.cross_entropy(z, y).item(), "by_class": [{"label": k, "n": int((y == k).sum()), "correct": int(correct[y == k].sum())} for k in range(10)]} @torch.no_grad() def evaluate(net, x, y, ids): logits = net(x[ids]).tolist() labels = y[ids].tolist() return {"ids": ids, "logits": logits, "labels": labels, "metrics": metrics(logits, labels)} 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"] observed = y.clone() for i, label in plan["observed_labels"].items(): observed[int(i)] = label start = time.monotonic() torch.manual_seed(plan["selector_seed"]) selector = model(plan) selector_history = train(selector, x, y, plan["selector_ids"], plan, plan["selector_seed"], plan["selector_steps"]) with torch.no_grad(): losses = F.cross_entropy(selector(x[plan["candidate_ids"]]), observed[plan["candidate_ids"]], reduction="none").tolist() scores = dict(zip(plan["candidate_ids"], losses)) record = {"plan_sha256": lab.digest(plan_path), "source_sha256": {p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]}, "environment": {"python": platform.python_version(), "torch": torch.__version__, "sklearn": sklearn.__version__, "numpy": np.__version__, "platform": platform.platform(), "device": "cpu", "threads": 1}, "parameters": sum(p.numel() for p in selector.parameters()), "selector_history": selector_history, "candidate_scores": scores, "candidate_true_labels": {str(i): int(y[i]) for i in plan["candidate_ids"]}, "runs": []} for seed in plan["seeds"]: torch.manual_seed(seed) initial = model(plan) for arm in plan["arms"]: selected = choose(plan, observed, scores, seed, arm) net = deepcopy(initial) history = train(net, x, observed, selected, plan, seed, plan["steps"], arm == "curriculum") evaluations = {split: evaluate(net, x, y, plan[f"{split}_ids"]) for split in ["dev", "test"]} path = out/f"seed-{seed}-{arm}.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, plan[f"{split}_ids"]) == ev for split, ev in evaluations.items()) counts = lambda labels: [int((labels[selected] == k).sum()) for k in range(10)] row = {"seed": seed, "arm": arm, "selected": selected, "true_class_counts": counts(y), "observed_class_counts": counts(observed), "selected_changed": sum(i in plan["changed_ids"] for i in selected), "history": history, "evaluations": evaluations, "checkpoint": {"file": path.name, "sha256": lab.digest(path), "reload_equal": True}} record["runs"].append(row) m = evaluations["test"]["metrics"] print(f"seed={seed} arm={arm} test={m['correct']}/{m['n']} " f"tail={m['tail_correct']}/{m['tail_n']}", flush=True) record["wall_seconds"] = time.monotonic()-start (out/"record.json").write_text(json.dumps(record, ensure_ascii=False, separators=(",", ":"))+"\n") print(f"完成 18 次学生训练：{out}", flush=True) 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: with args.freeze.open("x") as stream: json.dump(make_plan(), stream, ensure_ascii=False, indent=2) stream.write("\n") print(f"已冻结数据、噪声与训练设置：{args.freeze}") return if not args.run: print("默认不训练。请使用 --run --out-dir runs/data-efficiency-new。") return if args.out_dir is None: parser.error("需要 --out-dir，且目录不能已存在") run(args.plan, args.out_dir) if __name__ == "__main__": main()