| """ |
| lpft - whole-network adaptation of the DANN MLP with the shots: LP-FT (Kumar et al. |
| 2022) with an L2-SP anchor (Li et al. 2018), optionally plus SHOT-style information |
| maximisation on unlabelled Amsterdam rows (Liang et al. 2020). |
| |
| python tools/nnlab/lpft.py [--smoke] [--only <substring>] [--base-only] |
| |
| Base network: tab_nn.py's "dann" backbone on input "rz" (lab.robustz per city), seed 0, |
| trained exactly as tab_nn.train does (epochs from CACHE/tab_nn/selection.json), once, |
| with NNLAB_BASE_THREADS (default 3) torch threads - the count tab_nn's cache was made with; |
| CPU results depend on it (2 threads gave Za max abs diff 14) - then the run's NNLAB_THREADS; |
| state_dict cached at CACHE/lpft/dann_rz_s0.pt. Its Amsterdam embedding is compared with |
| CACHE/tab_nn/emb_dann_rz_s0.npz Za (max abs diff logged and stored in meta). |
| |
| Per protocol draw (fresh copy of the base each time): |
| LP last Linear <- lab.finetune_last_layer(W0, b0, Z_shot, y_shot, lam), lam from |
| selection.json (= tab_dann_rz_ft for seed 0 only; not saved). |
| FT everything unfrozen (domain head unused); BatchNorm in eval mode (running stats |
| frozen), dropout off (net.eval() throughout); AdamW lr 1e-4 backbone, 1e-3 last |
| layer, weight decay 0; STEPS=60 full-batch steps; loss = |
| class-balanced CE(shots) |
| + 0.1 * CE(fresh class-balanced Madrid replay batch, 128 per class = 512) |
| + 0.01 * sum_params ||theta - theta_LP||^2 |
| [+ 0.1 * (mean entropy - entropy of mean prediction) on 512 random Amsterdam |
| rows, variant "_ent"]. No early stopping. |
| Methods: tab_dann_rz_lpft, tab_dann_rz_lpft_ent (suffix _SMOKE under --smoke). |
| |
| Madrid check (evidence only, run once, CACHE/lpft/madrid_check.json, copied into meta): |
| lab.madrid_split; base trained on the source part with the unlabelled pseudo-target |
| as adversary (transforms per part, as tab_nn's selection); for n in (5, 50, 200) and |
| 3 draws (rng(9000 + 100 n + rep)) from the pseudo-target: LP vs LP-FT vs LP-FT-ent, |
| macro F1 of W.decide (source-part prior) on the remaining pseudo-target rows. |
| Replay batches come from the source part; the entropy batch from the pseudo-target. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import lab |
| import numpy as np |
| import tab_nn |
|
|
| W = lab.W |
| T0 = time.time() |
|
|
| KIND, INP, SEED = "dann", "rz", 0 |
| STEPS = 60 |
| LR_BACKBONE, LR_HEAD = 1e-4, 1e-3 |
| REPLAY_W, REPLAY_PER_CLASS = 0.1, 128 |
| L2SP_W = 0.01 |
| ENT_W, ENT_BATCH = 0.1, 512 |
| CHECK_NS, CHECK_REPS = (5, 50, 200), 3 |
| BASE_THREADS = int(os.environ.get("NNLAB_BASE_THREADS", "3")) |
| METHODS = {"lpft": False, "lpft_ent": True} |
|
|
|
|
| def log(msg): |
| print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True) |
|
|
|
|
| |
| |
| |
| def make_net(state, d_in): |
| net = tab_nn.build_net(KIND, d_in) |
| net.load_state_dict(state) |
| net.eval() |
| return net |
|
|
|
|
| def logits_of(net, X, bs=8192): |
| import torch |
| out = [] |
| with torch.no_grad(): |
| for k in range(0, len(X), bs): |
| out.append(net(torch.from_numpy(np.ascontiguousarray(X[k:k + bs], np.float32)))[0].double().numpy()) |
| return np.vstack(out) |
|
|
|
|
| def get_base(path, Xs, ys, Xt, epochs): |
| import torch |
| if path.exists(): |
| return torch.load(path), None |
| t0 = time.time() |
| |
| |
| |
| run_threads, lab.THREADS = lab.THREADS, BASE_THREADS |
| try: |
| net, _ = tab_nn.train(KIND, Xs, ys, Xt, epochs, SEED) |
| finally: |
| lab.THREADS = run_threads |
| torch.set_num_threads(run_threads) |
| state = {k: v.detach().clone() for k, v in net.state_dict().items()} |
| torch.save(state, path) |
| log(f"trained base {path.name} ({BASE_THREADS} threads): {epochs} epochs on {min(len(Xs), tab_nn.SMOKE_ROWS or len(Xs))} rows {time.time() - t0:.0f}s") |
| return state, net |
|
|
|
|
| |
| |
| |
| def lpft(state, X_shot, ys, X_pred, X_rep, y_rep, X_unl, lam, seed, variants): |
| """Returns {"lp": P, suffix: P ...} probabilities on X_pred (columns classes 1..4).""" |
| import torch |
| import torch.nn.functional as F |
| lab.torch_setup(seed) |
| base = make_net(state, X_shot.shape[1]) |
| Z_shot = tab_nn.embed(base, X_shot) |
| W0, b0 = tab_nn.last_layer(base) |
| Wf, bf = lab.finetune_last_layer(W0, b0, Z_shot, ys, lam) |
| with torch.no_grad(): |
| base.head.weight.copy_(torch.tensor(Wf, dtype=torch.float32)) |
| base.head.bias.copy_(torch.tensor(bf, dtype=torch.float32)) |
| out = {"lp": lab.softmax(logits_of(base, X_pred))} |
|
|
| Xs = torch.from_numpy(np.ascontiguousarray(X_shot, np.float32)) |
| Ys = torch.from_numpy(np.asarray(ys) - 1).long() |
| cw = torch.tensor(lab.balanced_class_weights(np.asarray(ys)), dtype=torch.float32) |
| Xr = torch.from_numpy(np.ascontiguousarray(X_rep, np.float32)) |
| by_cls = [torch.from_numpy(np.flatnonzero(np.asarray(y_rep) == c)) for c in lab.CLASSES] |
| Xu = torch.from_numpy(np.ascontiguousarray(X_unl, np.float32)) |
|
|
| for sfx, use_ent in variants.items(): |
| net = copy.deepcopy(base) |
| net.eval() |
| if net.dom is not None: |
| for p in net.dom.parameters(): |
| p.requires_grad_(False) |
| named = [(k, p) for k, p in net.named_parameters() if p.requires_grad] |
| anchor = {k: p.detach().clone() for k, p in named} |
| head_p = [p for k, p in named if k.startswith("head.")] |
| body_p = [p for k, p in named if not k.startswith("head.")] |
| opt = torch.optim.AdamW([{"params": body_p, "lr": LR_BACKBONE}, |
| {"params": head_p, "lr": LR_HEAD}], weight_decay=0.0) |
| g = torch.Generator().manual_seed(seed) |
| |
| |
| g_ent = torch.Generator().manual_seed(seed + 7_777_777) |
| for _ in range(STEPS): |
| loss = F.cross_entropy(net(Xs)[0], Ys, weight=cw) |
| ri = torch.cat([ix[torch.randint(0, len(ix), (REPLAY_PER_CLASS,), generator=g)] for ix in by_cls]) |
| yr = torch.repeat_interleave(torch.arange(4), REPLAY_PER_CLASS) |
| loss = loss + REPLAY_W * F.cross_entropy(net(Xr[ri])[0], yr) |
| loss = loss + L2SP_W * sum(((p - anchor[k]) ** 2).sum() for k, p in named) |
| if use_ent: |
| j = torch.randint(0, len(Xu), (ENT_BATCH,), generator=g_ent) |
| pu = F.softmax(net(Xu[j])[0], 1) |
| ment = -(pu * torch.log(pu + 1e-8)).sum(1).mean() |
| pm = pu.mean(0) |
| hm = -(pm * torch.log(pm + 1e-8)).sum() |
| loss = loss + ENT_W * (ment - hm) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| net.eval() |
| out[sfx] = lab.softmax(logits_of(net, X_pred)) |
| return out |
|
|
|
|
| |
| |
| |
| def madrid_check(d, cache, epochs, lam, smoke): |
| path = cache / "madrid_check.json" |
| if path.exists(): |
| return json.loads(path.read_text()) |
| src, tgt = lab.madrid_split(d["pm"]) |
| Xs, Xt = tab_nn.TRANSFORM[INP](d["Xm"][src]), tab_nn.TRANSFORM[INP](d["Xm"][tgt]) |
| ys, yt = d["ym"][src], d["ym"][tgt] |
| prior = np.array([(ys == c).mean() for c in lab.CLASSES]) |
| state, net = get_base(cache / f"{KIND}_{INP}_src_s{SEED}.pt", Xs, ys, Xt, epochs) |
| res = {"ns": list(CHECK_NS), "reps": CHECK_REPS, "shot_rng": "default_rng(9000 + 100 n + rep)", |
| "epochs": epochs, "lam": lam, "draws": []} |
| sel_f = lab.CACHE / "tab_nn" / f"sel_{KIND}.npz" |
| if net is not None and sel_f.exists() and not smoke: |
| with np.load(sel_f) as z: |
| diff = float(np.abs(tab_nn.embed(net, Xt) - z["Zt"]).max()) |
| res["src_base_vs_tab_nn_sel_Zt_maxabs"] = diff |
| log(f"madrid-check base vs tab_nn sel_{KIND}.npz Zt: max abs diff {diff:.2e}") |
| ns = (5, 200) if smoke else CHECK_NS |
| reps = 1 if smoke else CHECK_REPS |
| for n in ns: |
| for rep in range(reps): |
| t0 = time.time() |
| rng = np.random.default_rng(9000 + 100 * n + rep) |
| shot = np.concatenate([rng.choice(np.flatnonzero(yt == c), n, replace=False) for c in lab.CLASSES]) |
| keep = np.setdiff1d(np.arange(len(yt)), shot) |
| P = lpft(state, Xt[shot], yt[shot], Xt[keep], Xs, ys, Xt, lam, 10_000 + 1000 * n + rep, METHODS) |
| row = {"n": n, "rep": rep} |
| for k, Pk in P.items(): |
| row[k] = W.macro_f1(yt[keep], W.decide(Pk, prior)) |
| res["draws"].append(row) |
| log(f"madrid-check n={n} rep={rep}: " + " ".join(f"{k}={v:.4f}" for k, v in row.items() |
| if k not in ("n", "rep")) + f" {time.time() - t0:.1f}s") |
| import pandas as pd |
| df = pd.DataFrame(res["draws"]) |
| res["mean"] = {k: {int(n): float(v) for n, v in df.groupby("n")[k].mean().items()} |
| for k in ["lp", *METHODS]} |
| path.write_text(json.dumps(res, indent=1)) |
| log(f"madrid-check means: {res['mean']}") |
| return res |
|
|
|
|
| |
| |
| |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--smoke", action="store_true") |
| ap.add_argument("--only", default="") |
| ap.add_argument("--base-only", action="store_true", |
| help="train/cache the Amsterdam base net, compare with tab_nn's embedding, stop") |
| a = ap.parse_args() |
| smoke = a.smoke |
| sfx = "_SMOKE" if smoke else "" |
| if smoke: |
| tab_nn.SMOKE_ROWS = 3000 |
| global BASE_THREADS |
| BASE_THREADS = lab.THREADS |
| cache = lab.CACHE / ("lpft" + sfx) |
| cache.mkdir(parents=True, exist_ok=True) |
| sel = json.loads((lab.CACHE / "tab_nn" / "selection.json").read_text()) |
| epochs = 2 if smoke else int(sel["epochs"][KIND]) |
| lam = float(sel["lam"]) |
| draws = [(5, 0), (200, 0)] if smoke else list(lab.draws()) |
| names = {s: f"tab_{KIND}_{INP}_{s}{sfx}" for s in METHODS} |
| variants = {s: e for s, e in METHODS.items() if a.only in names[s]} |
|
|
| d = lab.load(official60=False) |
| Xm, Xa = tab_nn.city_inputs(d, INP) |
| log("data loaded") |
|
|
| |
| base_path = cache / f"{KIND}_{INP}_s{SEED}.pt" |
| state, net = get_base(base_path, Xm, d["ym"], Xa, epochs) |
| chk_path = cache / "base_check.json" |
| if net is not None or not chk_path.exists(): |
| if net is None: |
| net = make_net(state, Xm.shape[1]) |
| emb_f = lab.CACHE / "tab_nn" / f"emb_{KIND}_{INP}_s{SEED}.npz" |
| chk = {"compared_with": str(emb_f.relative_to(lab.ROOT)), "base_threads": BASE_THREADS} |
| if smoke: |
| chk["note"] = "smoke base (3000 rows, 2 epochs): no comparison" |
| else: |
| with np.load(emb_f) as z: |
| Za = tab_nn.embed(net, Xa) |
| chk["Za_maxabs"] = float(np.abs(Za - z["Za"]).max()) |
| chk["W0_maxabs"] = float(np.abs(tab_nn.last_layer(net)[0] - z["W0"]).max()) |
| chk["match_1e-3"] = bool(chk["Za_maxabs"] <= 1e-3) |
| assert int(z["epochs"]) == epochs |
| msg = "MATCH" if chk["match_1e-3"] else "MISMATCH (continuing; threading?)" |
| log(f"base vs tab_nn {emb_f.name}: Za max abs diff {chk['Za_maxabs']:.2e}, " |
| f"W0 {chk['W0_maxabs']:.2e} -> {msg}") |
| chk_path.write_text(json.dumps(chk, indent=1)) |
| base_chk = json.loads(chk_path.read_text()) |
| del net |
| if a.base_only: |
| log("base only: done") |
| return |
| if not variants: |
| log("nothing selected by --only") |
| return |
|
|
| check = madrid_check(d, cache, epochs, lam, smoke) |
|
|
| meta = dict(script="tools/nnlab/lpft.py", backbone=KIND, input=INP, seed=SEED, epochs=epochs, |
| lp_lam=lam, steps=STEPS, lr_backbone=LR_BACKBONE, lr_head=LR_HEAD, weight_decay=0.0, |
| replay_weight=REPLAY_W, replay_batch=4 * REPLAY_PER_CLASS, l2sp_weight=L2SP_W, |
| ent_weight=ENT_W, ent_batch=ENT_BATCH, bn="eval (running stats frozen)", dropout="off", |
| draw_seed="torch seed / generator 1000 n + rep", |
| base_threads=BASE_THREADS, base_check=base_chk, madrid_check=check, smoke=smoke, |
| hyperparameters="fixed in advance by the brief; epochs and lam from tab_nn Madrid selection") |
| for n, rep in draws: |
| todo = {s: e for s, e in variants.items() if not lab.have_pred(names[s], n, rep)} |
| if not todo: |
| continue |
| t0 = time.time() |
| shot = lab.shots_for(d, n, rep) |
| ys = lab.shot_labels(d, shot) |
| P = lpft(state, Xa[shot], ys, Xa, Xm, d["ym"], Xa, lam, 1000 * n + rep, todo) |
| for s in todo: |
| lab.save_pred(names[s], n, rep, P[s], {**meta, "method": names[s], "entropy_term": todo[s]}) |
| log(f"{'+'.join(names[s] for s in todo)} n={n} rep={rep} {time.time() - t0:.1f}s") |
| log("done") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
