| """ |
| tab_nn - tabular neural networks with a FROZEN backbone and last-layer adaptation. |
| |
| python tools/nnlab/tab_nn.py [--smoke] [--only <substring>] |
| |
| Backbones (trained ONCE on all Madrid labels, class-balanced CE, AdamW lr 1e-3, |
| wd 1e-4, batch 512, seeds 0,1,2): |
| mlp Linear-BN-ReLU-Drop(0.2) x2 (256) -> Linear(64)-ReLU -> Linear(64,64)-ReLU -> Linear(64,4) |
| res ResNet-MLP (Gorishniy et al. 2021): Linear(256), 3 x [x + Lin(Drop(ReLU(Lin(BN(x)))))], |
| BN-ReLU -> Linear(64)-ReLU -> Linear(64,4) |
| dann mlp + domain head (Lin(64,64)-ReLU-Lin(64,1)) on the first 64-d features through |
| gradient reversal, lambda = 2/(1+exp(-10p))-1, adversary vs UNLABELLED Amsterdam rows |
| (as tools/compare_teammate.py train_dann; Madrid and Amsterdam batches go through |
| the net separately) |
| coral mlp + deep-CORAL loss (weight 1.0) between the penultimate features of a Madrid |
| batch and an unlabelled Amsterdam batch |
| Inputs: rz = 393 feats lab.robustz per city; rg = 393 feats lab.rankgauss per city; |
| o60 = official 60 feats lab.standardise per city. |
| The EMBEDDING is the input of the final Linear(.,4); (W0, b0) are that layer's weights. |
| |
| Heads per protocol draw (probabilities averaged over the 3 seeds before saving): |
| zs source head unchanged (saved once, n=0 rep=0) |
| ft lab.finetune_last_layer(W0, b0, Z_shot, y_shot, lam) lam chosen on Madrid |
| joint the objective of lab.joint_head(Z_madrid, y_madrid, Z_shot, y_shot, shot_weight=30) |
| (same weights, C=1) on embeddings standardised with the Madrid embedding mean/sd |
| (same affine map for the shots), solved with newton-cholesky warm-started from |
| the Madrid-only fit (see head_joint: lab.joint_head's lbfgs does not converge |
| in 400 iterations and would cost ~2 h over the run) |
| proto prototypes (n_c * shot mean + k * Madrid class mean)/(n_c + k), k=10, embeddings |
| centred by their own city's mean and L2-normalised; P = softmax(-20 * sq. distance) |
| Method names: tab_{backbone}_{input}_{head}. "rg" is run for mlp and dann only (time budget), |
| so 10 backbone x input combos x 4 heads = 40 methods. |
| |
| Madrid-only choices (lab.madrid_split, pseudo-target = the 25 % of 2 km blocks furthest |
| along axis 0; transforms computed separately on the source part and on the pseudo-target, |
| exactly as they are per city; adversary / CORAL target = the unlabelled pseudo-target): |
| * epochs per backbone type from {15, 30, 45}: input rz, seed 0, trained on the source |
| part, scored by macro F1 of the decision rule (W.decide, source-part prior) on the |
| pseudo-target. Non-adversarial nets: one 45-epoch run evaluated after 15/30/45 epochs |
| (identical to separate runs: constant lr, same seed). dann: separate runs, since its |
| lambda schedule depends on the total epoch count. |
| * lam for ft, ONE value for all backbones and n, from {0.1, 1, 10, 100}: on each |
| backbone's selected source-part net (rz, seed 0), draw n in {5, 50} per class from the |
| pseudo-target (5 draws, rng(7000 + 100 n + rep)), fine-tune, score the remaining |
| pseudo-target with the rule; pick the lam with the best mean over backbones x n x draws. |
| Both are cached in CACHE/tab_nn/selection.json and copied into every meta.json. |
| |
| CACHE LAYOUT (lab.CACHE / "tab_nn", or "tab_nn_SMOKE" for --smoke): |
| selection.json epochs per backbone, lam, all pseudo-target scores |
| sel_{backbone}.npz source-part (madrid_split) net at the chosen epochs: |
| Zs, Zt (source / pseudo-target embeddings), W0, b0 |
| emb_{backbone}_{input}_s{seed}.npz one file per (backbone, input, seed); inputs rz,o60 |
| for all backbones, rg for mlp and dann; seeds 0,1,2: |
| Zm (76263, 64) float32 Madrid embeddings (feature-table row order, eval mode) |
| Za (25992, 64) float32 Amsterdam embeddings (feature-table row order) |
| W0 (4, 64) float32 final layer weight; b0 (4,) final layer bias |
| epochs () int epochs it was trained for |
| Logits = Z @ W0.T + b0, columns = classes 1..4. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import lab |
| import numpy as np |
|
|
| W = lab.W |
| T0 = time.time() |
|
|
| BACKBONES = ("mlp", "res", "dann", "coral") |
| INPUTS = ("rz", "rg", "o60") |
| |
| |
| INPUTS_FOR = {"mlp": INPUTS, "dann": INPUTS, "res": ("rz", "o60"), "coral": ("rz", "o60")} |
| HEADS = ("zs", "ft", "joint", "proto") |
| SEEDS = (0, 1, 2) |
| EPOCH_GRID = (15, 30, 45) |
| LAM_GRID = (0.1, 1.0, 10.0, 100.0) |
| LAM_NS = (5, 50) |
| LAM_REPS = 5 |
| SHOT_WEIGHT = 30.0 |
| PROTO_K, PROTO_T = 10.0, 20.0 |
| CORAL_W = 1.0 |
| BATCH, LR, WD = 512, 1e-3, 1e-4 |
|
|
|
|
| def log(msg): |
| print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True) |
|
|
|
|
| |
| |
| |
| def build_net(kind: str, d_in: int): |
| import torch.nn as nn |
| from torch.autograd import Function |
|
|
| class GradReverse(Function): |
| @staticmethod |
| def forward(ctx, x, lam): |
| ctx.lam = lam |
| return x.view_as(x) |
|
|
| @staticmethod |
| def backward(ctx, g): |
| return -ctx.lam * g, None |
|
|
| class ResBlock(nn.Module): |
| def __init__(self, h, p=0.2): |
| super().__init__() |
| self.f = nn.Sequential(nn.BatchNorm1d(h), nn.Linear(h, h), nn.ReLU(), nn.Dropout(p), nn.Linear(h, h)) |
|
|
| def forward(self, x): |
| return x + self.f(x) |
|
|
| class Net(nn.Module): |
| def __init__(self): |
| super().__init__() |
| if kind == "res": |
| self.f1 = nn.Sequential(nn.Linear(d_in, 256), ResBlock(256), ResBlock(256), ResBlock(256), |
| nn.BatchNorm1d(256), nn.ReLU(), nn.Linear(256, 64), nn.ReLU()) |
| self.f2 = nn.Identity() |
| else: |
| self.f1 = nn.Sequential(nn.Linear(d_in, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2), |
| nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2), |
| nn.Linear(256, 64), nn.ReLU()) |
| self.f2 = nn.Sequential(nn.Linear(64, 64), nn.ReLU()) |
| self.head = nn.Linear(64, 4) |
| self.dom = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 1)) if kind == "dann" else None |
|
|
| def forward(self, x, lam=0.0): |
| f = self.f1(x) |
| z = self.f2(f) |
| dl = self.dom(GradReverse.apply(f, lam)) if self.dom is not None else None |
| return self.head(z), z, dl |
|
|
| return Net() |
|
|
|
|
| def coral_loss(a, b): |
| import torch |
| d = a.shape[1] |
|
|
| def cov(x): |
| x = x - x.mean(0, keepdim=True) |
| return x.T @ x / max(len(x) - 1, 1) |
| return ((cov(a) - cov(b)) ** 2).sum() / (4 * d * d) |
|
|
|
|
| SMOKE_ROWS = None |
|
|
|
|
| def train(kind, Xs, ys, Xt, epochs, seed, eval_at=(), eval_fn=None): |
| """Train a backbone. Xt = unlabelled target rows (dann / coral). If eval_at is |
| given, eval_fn(net) is called after those epochs; returns (net, {ep: result}).""" |
| torch = lab.torch_setup(seed) |
| if SMOKE_ROWS and len(Xs) > SMOKE_ROWS: |
| sub = np.random.default_rng(seed).choice(len(Xs), SMOKE_ROWS, replace=False) |
| Xs, ys = Xs[sub], np.asarray(ys)[sub] |
| import torch.nn as nn |
| net = build_net(kind, Xs.shape[1]) |
| opt = torch.optim.AdamW(net.parameters(), lr=LR, weight_decay=WD) |
| cw = torch.tensor(lab.balanced_class_weights(ys), dtype=torch.float32) |
| ce, bce = nn.CrossEntropyLoss(weight=cw), nn.BCEWithLogitsLoss() |
| S, Y = torch.from_numpy(np.ascontiguousarray(Xs, np.float32)), torch.from_numpy(np.asarray(ys) - 1).long() |
| T = torch.from_numpy(np.ascontiguousarray(Xt, np.float32)) |
| g = torch.Generator().manual_seed(1000 + seed) |
| steps = int(np.ceil(len(S) / BATCH)) |
| res = {} |
| for ep in range(epochs): |
| net.train() |
| perm = torch.randperm(len(S), generator=g) |
| for k in range(steps): |
| i = perm[k * BATCH:(k + 1) * BATCH] |
| if len(i) < 2: |
| continue |
| if kind == "dann": |
| p = (ep * steps + k) / (epochs * steps) |
| lam = 2 / (1 + np.exp(-10 * p)) - 1 |
| j = torch.randint(0, len(T), (len(i),), generator=g) |
| ls, _, ds = net(S[i], lam) |
| _, _, dt = net(T[j], lam) |
| loss = ce(ls, Y[i]) + bce(ds.squeeze(1), torch.zeros(len(i))) + bce(dt.squeeze(1), torch.ones(len(j))) |
| elif kind == "coral": |
| j = torch.randint(0, len(T), (len(i),), generator=g) |
| ls, zs, _ = net(S[i]) |
| _, zt, _ = net(T[j]) |
| loss = ce(ls, Y[i]) + CORAL_W * coral_loss(zs, zt) |
| else: |
| ls, _, _ = net(S[i]) |
| loss = ce(ls, Y[i]) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| if ep + 1 in eval_at: |
| res[ep + 1] = eval_fn(net) |
| net.eval() |
| return net, res |
|
|
|
|
| def embed(net, X): |
| import torch |
| net.eval() |
| out = [] |
| with torch.no_grad(): |
| for k in range(0, len(X), 8192): |
| out.append(net(torch.from_numpy(np.ascontiguousarray(X[k:k + 8192], np.float32)))[1].numpy()) |
| return np.vstack(out).astype(np.float32) |
|
|
|
|
| def last_layer(net): |
| return (net.head.weight.detach().numpy().astype(np.float32).copy(), |
| net.head.bias.detach().numpy().astype(np.float32).copy()) |
|
|
|
|
| |
| |
| |
| def _by_columns(f, chunk=48): |
| """Apply a per-column lab transform in column blocks: identical output (every lab |
| transform is column-wise), but bounds the float64 / rank temporaries - lab.rankgauss on |
| the full 76k x 393 Madrid matrix alone peaks at ~1.5 GB extra (reviewer fix, RAM cap).""" |
| def g(X): |
| return np.hstack([f(X[:, k:k + chunk]) for k in range(0, X.shape[1], chunk)]) |
| return g |
|
|
|
|
| TRANSFORM = {"rz": _by_columns(lab.robustz), "rg": _by_columns(lab.rankgauss), "o60": lab.standardise} |
|
|
|
|
| def raw_inputs(d, inp): |
| return (d["Xm60"], d["Xa60"]) if inp == "o60" else (d["Xm"], d["Xa"]) |
|
|
|
|
| _INPUT_CACHE = {} |
|
|
|
|
| def city_inputs(d, inp): |
| """Per-city transformed inputs, kept in memory for the run (~150 MB per 393-feat input).""" |
| if inp not in _INPUT_CACHE: |
| Xm, Xa = raw_inputs(d, inp) |
| f = TRANSFORM[inp] |
| _INPUT_CACHE[inp] = (f(Xm), f(Xa)) |
| return _INPUT_CACHE[inp] |
|
|
|
|
| |
| |
| |
| def head_ft(e, Zs, ys, lam): |
| Wf, bf = lab.finetune_last_layer(e["W0"], e["b0"], Zs, ys, lam) |
| return lab.softmax(e["Za"].astype(np.float64) @ Wf.T + bf) |
|
|
|
|
| def joint_init(e): |
| """Standardise embeddings with the Madrid mean/sd and fit the Madrid-only |
| weighted LR once per seed (warm start for every draw's joint head).""" |
| from sklearn.linear_model import LogisticRegression |
| if "_Zm_std" not in e: |
| mu, sd = e["Zm"].mean(0), e["Zm"].std(0) |
| sd = np.where(sd > 1e-8, sd, 1.0) |
| e["_Zm_std"] = ((e["Zm"] - mu) / sd).astype(np.float64) |
| e["_Za_std"] = ((e["Za"] - mu) / sd).astype(np.float64) |
| w = lab.balanced_class_weights(e["ym"])[e["ym"] - 1] |
| e["_m0"] = LogisticRegression(C=1.0, solver="newton-cholesky", max_iter=100).fit( |
| e["_Zm_std"], e["ym"], sample_weight=w) |
| return e |
|
|
|
|
| def head_joint(e, shot, ys): |
| """Same objective as lab.joint_head(Zm, ym, Zs, ys, shot_weight=30, C=1) |
| (identical sample weights), solved with sklearn's newton-cholesky solver |
| warm-started from the Madrid-only fit: lab.joint_head's lbfgs stops at |
| max_iter=400 unconverged on these embeddings and takes ~9 s per fit; this |
| converges in 2-4 Newton steps (~2 s) and lands closer to the true optimum.""" |
| from sklearn.linear_model import LogisticRegression |
| e = joint_init(e) |
| Zm, Za = e["_Zm_std"], e["_Za_std"] |
| X = np.vstack([Zm, Za[shot]]) |
| y = np.r_[np.asarray(e["ym"]), np.asarray(ys)] |
| w = np.r_[np.ones(len(e["ym"])), np.full(len(ys), SHOT_WEIGHT)] |
| w = w * lab.balanced_class_weights(y)[y - 1] |
| m = LogisticRegression(C=1.0, solver="newton-cholesky", max_iter=100, warm_start=True) |
| m.coef_, m.intercept_ = e["_m0"].coef_.copy(), e["_m0"].intercept_.copy() |
| m.fit(X, y, sample_weight=w) |
| return m.predict_proba(Za) |
|
|
|
|
| def _unit(Z): |
| Z = Z - Z.mean(0) |
| return Z / np.clip(np.linalg.norm(Z, axis=1, keepdims=True), 1e-12, None) |
|
|
|
|
| def head_proto(e, shot, ys): |
| if "_Um" not in e: |
| e["_Um"], e["_Ua"] = _unit(e["Zm"].astype(np.float64)), _unit(e["Za"].astype(np.float64)) |
| Um, Ua = e["_Um"], e["_Ua"] |
| protos = [] |
| for c in lab.CLASSES: |
| mm = Um[e["ym"] == c].mean(0) |
| sc = Ua[shot[ys == c]] |
| nc = len(sc) |
| protos.append((nc * sc.mean(0) + PROTO_K * mm) / (nc + PROTO_K) if nc else mm) |
| Pr = np.stack(protos) |
| d2 = (Ua ** 2).sum(1)[:, None] - 2 * Ua @ Pr.T + (Pr ** 2).sum(1)[None] |
| return lab.softmax(-PROTO_T * d2) |
|
|
|
|
| |
| |
| |
| def selection(d, cache, smoke): |
| path = cache / "selection.json" |
| sel = json.loads(path.read_text()) if path.exists() else {"epochs": {}, "epoch_scores": {}} |
| grid = (1, 2) if smoke else EPOCH_GRID |
| src, tgt = lab.madrid_split(d["pm"]) |
| Xs, Xt = TRANSFORM["rz"](d["Xm"][src]), TRANSFORM["rz"](d["Xm"][tgt]) |
| ys, yt = d["ym"][src], d["ym"][tgt] |
| prior = np.array([(ys == c).mean() for c in lab.CLASSES]) |
|
|
| def f1(net): |
| P = lab.softmax(embed(net, Xt).astype(np.float64) @ last_layer(net)[0].T.astype(np.float64) |
| + last_layer(net)[1]) |
| return W.macro_f1(yt, W.decide(P, prior)) |
|
|
| emb = {} |
| for kind in BACKBONES: |
| ep_path = cache / f"sel_{kind}.npz" |
| if kind in sel["epochs"] and ep_path.exists(): |
| z = np.load(ep_path) |
| emb[kind] = dict(Zs=z["Zs"], Zt=z["Zt"], W0=z["W0"], b0=z["b0"]) |
| continue |
| t0 = time.time() |
| nets = {} |
| if kind == "dann": |
| scores = {} |
| for E in grid: |
| net, _ = train(kind, Xs, ys, Xt, E, 0) |
| scores[E], nets[E] = f1(net), (embed(net, Xs), embed(net, Xt), *last_layer(net)) |
| else: |
| snap = {} |
|
|
| def ev(net): |
| r = f1(net) |
| snap[len(snap)] = (embed(net, Xs), embed(net, Xt), *last_layer(net)) |
| net.train() |
| return r |
| _, scores = train(kind, Xs, ys, Xt, max(grid), 0, eval_at=grid, eval_fn=ev) |
| for k, E in enumerate(grid): |
| nets[E] = snap[k] |
| best = max(grid, key=lambda E: (scores[E], -E)) |
| sel["epochs"][kind] = best |
| sel["epoch_scores"][kind] = {str(k): v for k, v in scores.items()} |
| Zs_, Zt_, W0, b0 = nets[best] |
| np.savez(ep_path, Zs=Zs_, Zt=Zt_, W0=W0, b0=b0) |
| emb[kind] = dict(Zs=Zs_, Zt=Zt_, W0=W0, b0=b0) |
| path.write_text(json.dumps(sel, indent=1)) |
| log(f"select epochs {kind}: pseudo-target rule F1 {scores} -> {best} ({time.time() - t0:.0f}s)") |
|
|
| if "lam" not in sel: |
| t0 = time.time() |
| reps = 1 if smoke else LAM_REPS |
| tab = {str(l): [] for l in LAM_GRID} |
| per = {} |
| for kind in BACKBONES: |
| e = emb[kind] |
| for n in LAM_NS: |
| for rep in range(reps): |
| rng = np.random.default_rng(7000 + 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) |
| for l in LAM_GRID: |
| Wf, bf = lab.finetune_last_layer(e["W0"], e["b0"], e["Zt"][shot], yt[shot], l) |
| P = lab.softmax(e["Zt"][keep].astype(np.float64) @ Wf.T + bf) |
| s = W.macro_f1(yt[keep], W.decide(P, prior)) |
| tab[str(l)].append(s) |
| per.setdefault(f"{kind}_n{n}_lam{l}", []).append(s) |
| means = {l: float(np.mean(v)) for l, v in tab.items()} |
| sel["lam"] = float(max(LAM_GRID, key=lambda l: (means[str(l)], l))) |
| sel["lam_scores"] = means |
| sel["lam_scores_detail"] = {k: float(np.mean(v)) for k, v in per.items()} |
| path.write_text(json.dumps(sel, indent=1)) |
| log(f"select lam: {means} -> {sel['lam']} ({time.time() - t0:.0f}s)") |
| return sel |
|
|
|
|
| |
| |
| |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--smoke", action="store_true") |
| ap.add_argument("--only", default="") |
| ap.add_argument("--seeds", type=int, default=len(SEEDS), |
| help="use only the first k of seeds 0,1,2 (default 3; method names unchanged, " |
| "recorded in meta) - a time-saving fallback, not the planned run") |
| a = ap.parse_args() |
| smoke = a.smoke |
| global SMOKE_ROWS |
| SMOKE_ROWS = 3000 if smoke else None |
| sfx = "_SMOKE" if smoke else "" |
| cache = lab.CACHE / ("tab_nn" + sfx) |
| cache.mkdir(parents=True, exist_ok=True) |
| seeds = (0,) if smoke else SEEDS[:a.seeds] |
| draws = [(5, 0), (200, 0)] if smoke else list(lab.draws()) |
|
|
| d = lab.load(official60=True) |
| log("data loaded") |
|
|
| def mname(bb, inp, h): |
| return f"tab_{bb}_{inp}_{h}{sfx}" |
|
|
| combos = [(bb, inp) for bb in BACKBONES for inp in INPUTS_FOR[bb] |
| if any(a.only in mname(bb, inp, h) for h in HEADS)] |
| sel = selection(d, cache, smoke) |
| lam = sel["lam"] |
|
|
| for bb, inp in combos: |
| heads = [h for h in HEADS if a.only in mname(bb, inp, h)] |
| todo_draws = [(n, r) for n, r in draws |
| if any(not lab.have_pred(mname(bb, inp, h), n, r) for h in heads if h != "zs")] |
| need_zs = "zs" in heads and not lab.have_pred(mname(bb, inp, "zs"), 0, 0) |
| if not todo_draws and not need_zs: |
| log(f"{bb}/{inp}: all saved, skip") |
| continue |
| epochs = int(sel["epochs"][bb]) |
| embs, Xin = [], None |
| for s in seeds: |
| f = cache / f"emb_{bb}_{inp}_s{s}.npz" |
| if f.exists(): |
| with np.load(f) as z0: |
| stale = int(z0["epochs"]) != epochs |
| if stale: |
| log(f"cached {f.name} has a different epoch count than selection.json: retraining") |
| f.unlink() |
| if not f.exists(): |
| if Xin is None: |
| Xin = city_inputs(d, inp) |
| t0 = time.time() |
| net, _ = train(bb, Xin[0], d["ym"], Xin[1], epochs, s) |
| W0, b0 = last_layer(net) |
| np.savez(f, Zm=embed(net, Xin[0]), Za=embed(net, Xin[1]), W0=W0, b0=b0, epochs=epochs) |
| log(f"train {bb}/{inp} seed {s}: {epochs} epochs {time.time() - t0:.0f}s") |
| z = np.load(f) |
| embs.append(dict(Zm=z["Zm"], Za=z["Za"], W0=z["W0"], b0=z["b0"], ym=d["ym"])) |
| del Xin |
| meta = dict(script="tools/nnlab/tab_nn.py", backbone=bb, input=inp, seeds=list(seeds), |
| epochs=epochs, epoch_choice=sel["epoch_scores"].get(bb), ft_lam=lam, |
| lam_scores=sel.get("lam_scores"), joint_shot_weight=SHOT_WEIGHT, |
| proto_k=PROTO_K, proto_temp=PROTO_T, coral_weight=CORAL_W if bb == "coral" else None, |
| optim=dict(opt="AdamW", lr=LR, wd=WD, batch=BATCH), |
| joint_solver="newton-cholesky, warm start from Madrid-only fit, C=1", |
| selection="lab.madrid_split pseudo-target, rule macro F1, input rz seed 0", |
| smoke=smoke) |
| if need_zs: |
| P = np.mean([lab.softmax(e["Za"].astype(np.float64) @ e["W0"].T + e["b0"]) for e in embs], 0) |
| lab.save_pred(mname(bb, inp, "zs"), 0, 0, P, {**meta, "head": "zs"}) |
| log(f"{mname(bb, inp, 'zs')} n=0 saved") |
| for n, rep in todo_draws: |
| t0 = time.time() |
| shot = lab.shots_for(d, n, rep) |
| ys = lab.shot_labels(d, shot) |
| for h in heads: |
| if h == "zs" or lab.have_pred(mname(bb, inp, h), n, rep): |
| continue |
| th = time.time() |
| if h == "ft": |
| P = np.mean([head_ft(e, e["Za"][shot], ys, lam) for e in embs], 0) |
| elif h == "joint": |
| P = np.mean([head_joint(e, shot, ys) for e in embs], 0) |
| else: |
| P = np.mean([head_proto(e, shot, ys) for e in embs], 0) |
| lab.save_pred(mname(bb, inp, h), n, rep, P, {**meta, "head": h}) |
| log(f"{mname(bb, inp, h)} n={n} rep={rep} {time.time() - th:.1f}s") |
| del embs |
| log("done") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
