| """Sequence networks on the RAW yearly Landsat time series, frozen backbone + |
| last-layer adaptation. |
| |
| python tools/nnlab/ts_nn.py [--smoke] [--only <substring>] |
| |
| Step 1 tensor cache lab.CACHE/ts/{city}_T.npy [N, 42, 6] float32 DN, median of |
| QA-valid observations per year (src/features.py make_flat_view(agg="median"): |
| a year is kept only if every band has a valid median and Blue <= BLUE_MAX), |
| NaN where no kept observation; rows aligned with the processed feature table. |
| NOTE: the raw data covers 1984-2025 (42 years), as the feature pipeline uses. |
| Step 2 inputs 6 SR bands + NDVI, NDBI, MNDWI + observed mask = 10 ch x 42 steps, |
| linear interpolation in time (edges nearest, as to_year_matrices), indices from |
| the interpolated SR bands; per-city standardisation per channel. Variant "ry" |
| subtracts the city's per-year median over pixels (9 continuous channels) first. |
| Step 3 backbones tcnn, gru, tcnn_dann; class-balanced CE, AdamW(1e-3, wd 1e-4), |
| batch 256, constant LR, brightness/offset augmentation on the 6 band channels. |
| CPU budget: per (backbone, variant) ONE run on the lab.madrid_split source part, |
| checkpoints at EPOCH_GRID = (5, 10, 15), the checkpoint with the best |
| pseudo-target macro F1 IS the backbone (no retrain on all Madrid); lam for "ft" |
| chosen on the same model. 1 seed. Default combos: tcnn_st, tcnn_ry, |
| tcnn_dann_st; --extended adds tcnn_dann_ry, gru_st, gru_ry (~2.5 h more). |
| Step 4 heads zs (n=0), ft (lab.finetune_last_layer), joint (lab.joint_head); |
| probabilities averaged over seeds. |
| Methods: ts_{backbone}_{variant}_{head}. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| import json |
| import sys |
| import time |
| import warnings |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import lab |
| import numpy as np |
| import config as C |
|
|
| T0 = time.time() |
| YEARS = np.arange(1984, 2026) |
| NY = len(YEARS) |
| BANDS = list(C.BASE_BANDS) |
| TS_DIR = lab.CACHE / "ts" |
| NN_DIR = lab.CACHE / "ts_nn" |
| for _d in (TS_DIR, NN_DIR): |
| _d.mkdir(parents=True, exist_ok=True) |
|
|
| HEADS = ("zs", "ft", "joint") |
| EPOCH_GRID = (5, 10, 15) |
| LAM_GRID = (0.1, 1.0, 10.0, 100.0) |
| SEEDS = (0,) |
| DEFAULT_COMBOS = (("tcnn", "st"), ("tcnn", "ry"), ("tcnn_dann", "st")) |
| EXTENDED_COMBOS = (("tcnn_dann", "ry"), ("gru", "st"), ("gru", "ry")) |
| BATCH = 256 |
| LR, WD = 1e-3, 1e-4 |
| SHOT_WEIGHT = 30.0 |
|
|
|
|
| def _rss(): |
| try: |
| import psutil |
| mi = psutil.Process().memory_info() |
| return f" rss={mi.rss / 2**30:.2f}G peak={getattr(mi, 'peak_wset', mi.rss) / 2**30:.2f}G" |
| except Exception: |
| return "" |
|
|
|
|
| def log(msg): |
| print(f"[{time.time() - T0:7.1f}s] {msg}{_rss()}", flush=True) |
|
|
|
|
| |
| |
| |
| RAW = {"madrid": lab.ROOT / "data" / "raw" / "madrid_train.parquet", |
| "amsterdam": lab.ROOT / "data" / "raw" / "amsterdam_data.parquet"} |
|
|
|
|
| def build_tensor(city: str, P: np.ndarray) -> np.ndarray: |
| """P: (N, 2) px_key, py_key of the feature table (row order).""" |
| import pyarrow.parquet as pq |
| path = TS_DIR / f"{city}_T.npy" |
| if path.exists(): |
| T = np.load(path) |
| assert T.shape == (len(P), NY, 6) |
| return T |
| t0 = time.time() |
| cols = ["px_key", "py_key", "year"] + [f"{b}_{k}" for k in (1, 2, 3) for b in BANDS] \ |
| + [f"qa_valid_{k}" for k in (1, 2, 3)] |
| key_f = P[:, 0].astype(np.int64) * 10_000_000 + P[:, 1].astype(np.int64) |
| assert P[:, 1].min() >= 0 and P[:, 1].max() < 10_000_000 |
| order = np.argsort(key_f) |
| ks = key_f[order] |
| assert (np.diff(ks) > 0).all(), "duplicate pixel keys in feature table" |
| T = np.full((len(P), NY, 6), np.nan, np.float32) |
| seen = np.zeros((len(P), NY), bool) |
| f = pq.ParquetFile(RAW[city]) |
| n_rows = n_kept = 0 |
| for batch in f.iter_batches(batch_size=200_000, columns=cols): |
| tb = {c: batch.column(c) for c in cols} |
| n = batch.num_rows |
| n_rows += n |
| oks = [np.asarray(tb[f"qa_valid_{k}"].fill_null(False).to_numpy(zero_copy_only=False), bool) |
| for k in (1, 2, 3)] |
| med = np.empty((n, 6), np.float32) |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore", RuntimeWarning) |
| for j, b in enumerate(BANDS): |
| st = [] |
| for k in (1, 2, 3): |
| v = np.asarray(tb[f"{b}_{k}"].to_numpy(zero_copy_only=False), np.float32).copy() |
| v[~oks[k - 1]] = np.nan |
| st.append(v) |
| med[:, j] = np.nanmedian(np.stack(st), axis=0) |
| good = ~np.isnan(med).any(1) |
| good &= ~(med[:, 0] > C.BLUE_MAX) |
| px = np.asarray(tb["px_key"].to_numpy(zero_copy_only=False), np.int64) |
| py = np.asarray(tb["py_key"].to_numpy(zero_copy_only=False), np.int64) |
| yr = np.asarray(tb["year"].to_numpy(zero_copy_only=False), np.int64) |
| key = px * 10_000_000 + py |
| pos = np.searchsorted(ks, key).clip(0, len(ks) - 1) |
| hit = (ks[pos] == key) & good |
| assert ((yr[hit] >= YEARS[0]) & (yr[hit] <= YEARS[-1])).all() |
| r, c = order[pos[hit]], yr[hit] - YEARS[0] |
| assert not seen[r, c].any(), "duplicate pixel-year" |
| seen[r, c] = True |
| T[r, c] = med[hit] |
| n_kept += int(hit.sum()) |
| del batch, tb, med, st |
| found = seen.any(1) |
| assert found.all(), f"{(~found).sum()} feature-table pixels have no observation" |
| np.save(path, T) |
| log(f"tensor {city}: {n_rows:,} raw rows, {n_kept:,} pixel-years kept, " |
| f"shape {T.shape}, observed {seen.mean():.3f}, {time.time() - t0:.0f}s") |
| return T |
|
|
|
|
| |
| |
| |
| def interp_time(T: np.ndarray) -> np.ndarray: |
| """Linear interpolation over the year axis, nearest at the edges (per pixel, band).""" |
| N, Y, B = T.shape |
| out = T.copy() |
| x = np.arange(Y, dtype=np.float32) |
| for b in range(B): |
| M = out[:, :, b] |
| obs = ~np.isnan(M) |
| |
| idx = np.where(obs, np.arange(Y)[None, :], -1) |
| prev = np.maximum.accumulate(idx, axis=1) |
| idx2 = np.where(obs, np.arange(Y)[None, :], Y) |
| nxt = np.minimum.accumulate(idx2[:, ::-1], axis=1)[:, ::-1] |
| rows = np.arange(N)[:, None] |
| pv = np.where(prev >= 0, M[rows, prev.clip(0)], np.nan) |
| nv = np.where(nxt < Y, M[rows, nxt.clip(None, Y - 1)], np.nan) |
| w = np.where((prev >= 0) & (nxt < Y) & (nxt > prev), |
| (x[None, :] - prev) / np.maximum(nxt - prev, 1), 0.0).astype(np.float32) |
| lin = pv + w * (nv - pv) |
| fill = np.where(np.isnan(pv), nv, np.where(np.isnan(nv), pv, lin)) |
| out[:, :, b] = np.where(obs, M, fill) |
| assert not np.isnan(out).any() |
| return out |
|
|
|
|
| def make_inputs(T: np.ndarray, variant: str) -> np.ndarray: |
| """[N, Y, 6] DN -> [N, 10, Y] float32, standardised with this city's statistics.""" |
| obs = (~np.isnan(T).any(2)).astype(np.float32) |
| S = interp_time(T) * np.float32(C.SR_SCALE) + np.float32(C.SR_OFFSET) |
| B, G, R, N_, S1 = S[..., 0], S[..., 1], S[..., 2], S[..., 3], S[..., 4] |
| eps = np.float32(1e-6) |
| ndvi = (N_ - R) / (N_ + R + eps) |
| ndbi = (S1 - N_) / (S1 + N_ + eps) |
| mndwi = (G - S1) / (G + S1 + eps) |
| X = np.concatenate([S, ndvi[..., None], ndbi[..., None], mndwi[..., None], obs[..., None]], 2) |
| del S |
| X = np.clip(X, -5, 5) |
| if variant == "ry": |
| X[..., :9] -= np.median(X[..., :9], axis=0, keepdims=True) |
| mu = X.reshape(-1, 10).mean(0, dtype=np.float64) |
| sd = X.reshape(-1, 10).std(0, dtype=np.float64) |
| sd[sd < 1e-8] = 1 |
| X = ((X - mu.astype(np.float32)) / sd.astype(np.float32)).astype(np.float32) |
| return np.ascontiguousarray(X.transpose(0, 2, 1)) |
|
|
|
|
| |
| |
| |
| def build_net(kind: str): |
| import torch |
| nn = torch.nn |
|
|
| class GradReverse(torch.autograd.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 TempCNN(nn.Module): |
| def __init__(self, c_in=10, T=NY, dann=False): |
| super().__init__() |
| layers, c = [], c_in |
| for _ in range(3): |
| layers += [nn.Conv1d(c, 64, 5, padding=2), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.2)] |
| c = 64 |
| self.conv = nn.Sequential(*layers) |
| self.emb = nn.Sequential(nn.Flatten(), nn.Linear(64 * T, 128), nn.ReLU(), |
| nn.Linear(128, 64), nn.ReLU()) |
| self.head = nn.Linear(64, 4) |
| self.dom = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 1)) if dann else None |
|
|
| def embed(self, x): |
| return self.emb(self.conv(x)) |
|
|
| def forward(self, x, lam=0.0): |
| z = self.embed(x) |
| dom = self.dom(GradReverse.apply(z, lam)) if self.dom is not None else None |
| return self.head(z), dom |
|
|
| class GRUNet(nn.Module): |
| def __init__(self, c_in=10): |
| super().__init__() |
| self.gru = nn.GRU(c_in, 64, num_layers=2, bidirectional=True, batch_first=True) |
| self.emb = nn.Sequential(nn.Linear(128, 64), nn.ReLU()) |
| self.head = nn.Linear(64, 4) |
| self.dom = None |
|
|
| def embed(self, x): |
| _, h = self.gru(x.transpose(1, 2)) |
| return self.emb(torch.cat([h[-2], h[-1]], 1)) |
|
|
| def forward(self, x, lam=0.0): |
| return self.head(self.embed(x)), None |
|
|
| if kind == "tcnn": |
| return TempCNN() |
| if kind == "tcnn_dann": |
| return TempCNN(dann=True) |
| if kind == "gru": |
| return GRUNet() |
| raise ValueError(kind) |
|
|
|
|
| def augment(torch, xb, gen): |
| """Per-sample brightness scale U(0.9,1.1) and offset N(0,0.1) on the 6 band channels.""" |
| B = xb.shape[0] |
| s = torch.empty(B, 1, 1).uniform_(0.9, 1.1, generator=gen) |
| o = torch.randn(B, 1, 1, generator=gen) * 0.1 |
| xb = xb.clone() |
| xb[:, :6] = xb[:, :6] * s + o |
| return xb |
|
|
|
|
| def train(kind, Xs, ys, Xt, epochs, seed, eval_fn=None, eval_at=(), max_rows=None, sched_epochs=None): |
| """Returns (net, {epoch: eval result}, {epoch: state_dict copy}, sec/epoch). |
| sched_epochs: length of the DANN lambda schedule (default `epochs`); a retrain to a |
| selected checkpoint passes the selection run's length so it reproduces that model.""" |
| sched = sched_epochs or epochs |
| torch = lab.torch_setup(seed) |
| gen = torch.Generator().manual_seed(seed) |
| net = build_net(kind) |
| opt = torch.optim.AdamW(net.parameters(), lr=LR, weight_decay=WD) |
| counts = np.bincount(ys, minlength=5)[1:] |
| w = torch.tensor(len(ys) / (4 * np.maximum(counts, 1)), dtype=torch.float32) |
| ce, bce = torch.nn.CrossEntropyLoss(weight=w), torch.nn.BCEWithLogitsLoss() |
| S, Y = torch.from_numpy(Xs), torch.from_numpy(ys - 1).long() |
| Tt = torch.from_numpy(Xt) if Xt is not None else None |
| dann = kind.endswith("_dann") |
| n = len(S) if max_rows is None else min(len(S), max_rows) |
| steps = int(np.ceil(n / BATCH)) |
| evals, states, t_ep = {}, {}, [] |
| for ep in range(1, epochs + 1): |
| te = time.time() |
| net.train() |
| perm = torch.randperm(len(S), generator=gen)[:n] |
| for k in range(steps): |
| i = perm[k * BATCH:(k + 1) * BATCH] |
| if len(i) < 2: |
| continue |
| xb = augment(torch, S[i], gen) |
| if dann: |
| p = ((ep - 1) * steps + k) / (sched * steps) |
| lam = float(2 / (1 + np.exp(-10 * p)) - 1) |
| logit, dom_s = net(xb, lam) |
| j = torch.randint(0, len(Tt), (len(i),), generator=gen) |
| _, dom_t = net(augment(torch, Tt[j], gen), lam) |
| loss = ce(logit, Y[i]) + bce(dom_s.squeeze(1), torch.zeros(len(i))) + bce(dom_t.squeeze(1), torch.ones(len(i))) |
| else: |
| logit, _ = net(xb, 0.0) |
| loss = ce(logit, Y[i]) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| t_ep.append((time.time() - te) * len(S) / n) |
| if ep in eval_at: |
| states[ep] = copy.deepcopy(net.state_dict()) |
| if eval_fn is not None: |
| evals[ep] = eval_fn(net) |
| log(f" {kind} seed{seed} epoch {ep}/{epochs} {t_ep[-1]:.1f}s/full-epoch" |
| + (f" eval={evals[ep]:.4f}" if ep in evals else "")) |
| net.eval() |
| return net, evals, states, float(np.mean(t_ep)) |
|
|
|
|
| def embed(net, X, bs=2048): |
| import torch |
| net.eval() |
| out = [] |
| with torch.no_grad(): |
| for k in range(0, len(X), bs): |
| out.append(net.embed(torch.from_numpy(X[k:k + bs])).numpy()) |
| return np.concatenate(out).astype(np.float32) |
|
|
|
|
| def last_layer(net): |
| return (net.head.weight.detach().numpy().astype(np.float64).copy(), |
| net.head.bias.detach().numpy().astype(np.float64).copy()) |
|
|
|
|
| def proba(Z, Wt, b): |
| return lab.softmax(Z.astype(np.float64) @ Wt.T + b) |
|
|
|
|
| def mf1(y, P): |
| return lab.W.macro_f1(y, lab.CLASSES[P.argmax(1)]) |
|
|
| |
| |
| |
| |
| |
| def simulate_lam(Zt, yt, W0, b0, reps): |
| lam_score = {} |
| for lam in LAM_GRID: |
| sc = [] |
| for n in (5, 50): |
| for rep in range(reps): |
| rng = np.random.default_rng(7_000 + 1000 * n + rep) |
| sh = np.concatenate([rng.choice(np.flatnonzero(yt == c), min(n, int((yt == c).sum())), |
| replace=False) for c in lab.CLASSES]) |
| rest = np.setdiff1d(np.arange(len(yt)), sh) |
| Wf, bf = lab.finetune_last_layer(W0, b0, Zt[sh], yt[sh], lam) |
| sc.append(mf1(yt[rest], proba(Zt[rest], Wf, bf))) |
| lam_score[lam] = float(np.mean(sc)) |
| return lam_score |
|
|
|
|
| def backbone(kind, var, d, Xm, Xa, cfg): |
| """Returns (selection dict, [(Zm, Za, W0, b0) per seed]).""" |
| tag = f"{kind}_{var}{cfg['tag']}" |
| sel_path = NN_DIR / f"select_{tag}.json" |
| src, tgt = lab.madrid_split(d["pm"]) |
| ym = d["ym"] |
| grid = cfg["epoch_grid"] |
| embs = [] |
| sel = json.loads(sel_path.read_text()) if sel_path.exists() else None |
| for seed in cfg["seeds"]: |
| path = NN_DIR / f"{tag}_s{seed}.npz" |
| if path.exists() and sel is not None: |
| z = np.load(path) |
| embs.append((z["Zm"], z["Za"], z["W0"], z["b0"])) |
| continue |
| t0 = time.time() |
| if sel is None: |
| def ev(net): |
| Wt, b = last_layer(net) |
| return mf1(ym[tgt], proba(embed(net, Xm[tgt]), Wt, b)) |
| net, evals, states, sec = train(kind, Xm[src], ym[src], Xa, max(grid), seed, eval_fn=ev, |
| eval_at=grid, max_rows=cfg["max_rows"]) |
| E = max(grid, key=lambda e: (evals[e], -e)) |
| net.load_state_dict(states[E]) |
| net.eval() |
| W0, b0 = last_layer(net) |
| lam_score = simulate_lam(embed(net, Xm[tgt]), ym[tgt], W0, b0, cfg["lam_reps"]) |
| lam = max(LAM_GRID, key=lambda l: (lam_score[l], l)) |
| sel = {"backbone": kind, "variant": var, "seed_used": int(seed), |
| "split": "lab.madrid_split(pm) defaults (frac .25, axis 0, high)", |
| "epoch_grid": list(grid), |
| "pseudo_target_f1_by_epoch": {str(k): float(v) for k, v in evals.items()}, |
| "epochs": int(E), "lam_grid": list(LAM_GRID), |
| "pseudo_target_ft_f1_by_lam": {str(k): v for k, v in lam_score.items()}, |
| "lam": float(lam), |
| "lam_sim": f"n per class in (5, 50) x {cfg['lam_reps']} reps on the pseudo-target, " |
| "argmax macro F1 on the rest", |
| "sec_per_full_epoch": sec, "n_src": int(len(src)), "n_pseudo_target": int(len(tgt))} |
| sel_path.write_text(json.dumps(sel, indent=1)) |
| scores = {k: round(v, 4) for k, v in lam_score.items()} |
| log(f"select {tag}: epochs={E} {sel['pseudo_target_f1_by_epoch']} lam={lam} {scores}") |
| else: |
| net, _, _, sec = train(kind, Xm[src], ym[src], Xa, sel["epochs"], seed, |
| eval_at=(sel["epochs"],), max_rows=cfg["max_rows"], |
| sched_epochs=max(sel["epoch_grid"])) |
| Zm, Za = embed(net, Xm), embed(net, Xa) |
| W0, b0 = last_layer(net) |
| np.savez(path, Zm=Zm, Za=Za, W0=W0, b0=b0) |
| embs.append((Zm, Za, W0, b0)) |
| log(f"backbone {tag} seed{seed}: {sec:.1f}s/full-epoch, {time.time() - t0:.0f}s") |
| return sel, embs |
|
|
|
|
| def gru_check(): |
| """Smoke-only: the GRU network trains and embeds (the rest of its path is shared).""" |
| rng = np.random.default_rng(0) |
| X = rng.standard_normal((1024, 10, NY)).astype(np.float32) |
| y = rng.integers(1, 5, 1024) |
| net, _, _, sec = train("gru", X, y, None, 1, 0, eval_at=(1,)) |
| Z = embed(net, X[:300]) |
| assert Z.shape == (300, 64) and np.isfinite(Z).all() |
| log(f"gru check ok: ~{sec * 61932 / 1024:.0f}s per Madrid-source epoch (extrapolated)") |
|
|
|
|
| |
| |
| |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--smoke", action="store_true") |
| ap.add_argument("--only", default="") |
| ap.add_argument("--extended", action="store_true", |
| help="also tcnn_dann_ry, gru_st, gru_ry (beyond the 70 min budget)") |
| a = ap.parse_args() |
| combos = list(DEFAULT_COMBOS) + (list(EXTENDED_COMBOS) if a.extended else []) |
| if a.smoke: |
| cfg = dict(tag="_SMOKE", suffix="_SMOKE", epoch_grid=(1, 2), max_rows=4096, lam_reps=1, |
| seeds=(0,), draws=[(5, 0), (200, 0)]) |
| else: |
| cfg = dict(tag="", suffix="", epoch_grid=EPOCH_GRID, max_rows=None, lam_reps=5, |
| seeds=SEEDS, draws=list(lab.draws())) |
|
|
| def mname(bb, var, head): |
| return f"ts_{bb}_{var}_{head}{cfg['suffix']}" |
|
|
| todo = {(bb, var): [h for h in HEADS if a.only in mname(bb, var, h)] for bb, var in combos} |
| todo = {k: v for k, v in todo.items() if v} |
| log(f"methods: {[mname(bb, var, h) for (bb, var), hs in todo.items() for h in hs]}") |
|
|
| d = lab.load(official60=False) |
| Tm = build_tensor("madrid", d["pm"]) |
| Ta = build_tensor("amsterdam", d["pa"]) |
| if a.smoke: |
| gru_check() |
|
|
| for (bb, var), heads in todo.items(): |
| need = [h for h in heads for (n, rep) in ([(0, 0)] if h == "zs" else cfg["draws"]) |
| if not lab.have_pred(mname(bb, var, h), n, rep)] |
| if not need: |
| log(f"{bb}_{var}: all predictions present, skip") |
| continue |
| Xm, Xa = make_inputs(Tm, var), make_inputs(Ta, var) |
| log(f"inputs {var}: Madrid {Xm.shape} Amsterdam {Xa.shape}") |
| sel, embs = backbone(bb, var, d, Xm, Xa, cfg) |
| del Xm, Xa |
| E, lam = sel["epochs"], sel["lam"] |
| meta = {"family": "ts_nn", "backbone": bb, "variant": var, "seeds": list(cfg["seeds"]), |
| "epochs": E, "lam": lam, "shot_weight": SHOT_WEIGHT, "selection": sel, |
| "backbone_training_rows": "Madrid spatial source split only (lab.madrid_split)", |
| "joint_input": "embeddings standardised with Madrid embedding mean/std", |
| "input": f"10 ch x {NY} years (1984-2025), median of QA-valid obs, Blue<=BLUE_MAX", |
| "augment": "band scale U(0.9,1.1), offset N(0,0.1)", |
| "train": f"AdamW lr={LR} wd={WD} batch={BATCH} class-balanced CE, constant LR", |
| "smoke": bool(a.smoke)} |
| if "zs" in heads and not lab.have_pred(mname(bb, var, "zs"), 0, 0): |
| P = np.mean([proba(Za, W0, b0) for _, Za, W0, b0 in embs], 0) |
| lab.save_pred(mname(bb, var, "zs"), 0, 0, P, {**meta, "head": "zs"}) |
| log(f"{mname(bb, var, 'zs')} saved") |
| scaled = [] |
| for Zm, Za, _, _ in embs: |
| mu, sd = Zm.mean(0), Zm.std(0) + 1e-6 |
| scaled.append(((Zm - mu) / sd, (Za - mu) / sd)) |
| for n, rep in cfg["draws"]: |
| hs = [h for h in heads if h != "zs" and not lab.have_pred(mname(bb, var, h), n, rep)] |
| if not hs: |
| continue |
| t0 = time.time() |
| shot = lab.shots_for(d, n, rep) |
| ys = lab.shot_labels(d, shot) |
| for h in hs: |
| Ps = [] |
| for (Zm, Za, W0, b0), (Sm, Sa) in zip(embs, scaled): |
| if h == "ft": |
| Wf, bf = lab.finetune_last_layer(W0, b0, Za[shot], ys, lam) |
| Ps.append(proba(Za, Wf, bf)) |
| else: |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| m = lab.joint_head(Sm, d["ym"], Sa[shot], ys, shot_weight=SHOT_WEIGHT) |
| Ps.append(m.predict_proba(Sa)) |
| lab.save_pred(mname(bb, var, h), n, rep, np.mean(Ps, 0), {**meta, "head": h}) |
| log(f"{bb}_{var} n={n} rep={rep} heads={hs} {time.time() - t0:.1f}s") |
| log("done") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
