Buckets:
| #!/usr/bin/env python3 | |
| """Train GRU / Transformer / ConvNeXt on the dynamic sparse cache, GPU-resident. | |
| The whole sparse cache lives on the GPU; each batch is densified on-GPU (scatter), so there | |
| is zero per-batch CPU<->GPU transfer and no disk reads after load. Mirrors modeling/training.py | |
| train_one hyperparameters exactly (AdamW lr/wd, CE ignore PAD, grad clip 1.0, 15 epochs, | |
| best-validation-macro-F1 checkpoint selection). | |
| """ | |
| from __future__ import annotations | |
| import json, sys, time | |
| from pathlib import Path | |
| import numpy as np, torch, torch.nn as nn | |
| from sklearn.metrics import f1_score | |
| sys.path.insert(0, "/workspace/scope/modeling") | |
| from models import PAD, TOPICS, build_model, seed_everything, metrics, flatten_labeled | |
| ROOT = Path("/workspace/scope") | |
| CACHE = ROOT / "cache/dyn150_sparse/sparse.pt" | |
| DEVICE = "cuda" | |
| EPOCHS = 15 | |
| ARCHS = { | |
| "gru": dict(arch="gru", out="results/dyn150_2k/gru/gru_lr3e4", | |
| model={"hidden_width":256,"depth":3,"dropout":0.1}, lr=3e-4, bs=16), | |
| "transformer": dict(arch="transformer", out="results/dyn150_2k/transformer/transformer_wide", | |
| model={"hidden_width":256,"heads":8,"depth":2,"feedforward_width":512,"dropout":0.1}, lr=3e-4, bs=4), | |
| "convnext": dict(arch="convnext", out="results/dyn150_2k/tcn/convnext", | |
| model={"channels":256,"depth":6,"kernel":7,"expansion":4,"dropout":0.1,"layer_scale":0.01}, lr=1e-3, bs=16), | |
| } | |
| def load_cache(): | |
| d = torch.load(CACHE, map_location="cpu") | |
| W = int(d["width"]) | |
| col = d["col"].to(DEVICE); val = d["val"].to(DEVICE) | |
| tok_ptr = d["tok_ptr"].to(DEVICE) | |
| labels = d["labels"].to(DEVICE) | |
| lengths = d["lengths"].numpy() | |
| offsets = np.concatenate(([0], np.cumsum(lengths))) | |
| splits = {k: d[k].numpy() for k in ("train","validation","test")} | |
| print(f"resident: nnz={col.numel()} width={W} " | |
| f"gpu={torch.cuda.memory_allocated()/2**30:.2f} GiB", flush=True) | |
| return dict(col=col, val=val, tok_ptr=tok_ptr, labels=labels, | |
| offsets=offsets, splits=splits, W=W) | |
| def densify(C, rec_ids): | |
| """Build (B, maxT, W) fp32 features and (B, maxT) labels for a list of record ids, on GPU. | |
| Per-nnz token-local row indices are computed on the fly from tok_ptr (no giant resident | |
| row tensor).""" | |
| off = C["offsets"]; tok_ptr = C["tok_ptr"]; col = C["col"]; val = C["val"]; labels = C["labels"] | |
| Ts = [int(off[r+1]-off[r]) for r in rec_ids] | |
| maxT = max(Ts); B = len(rec_ids) | |
| x = torch.zeros(B, maxT, C["W"], device=DEVICE) | |
| y = torch.full((B, maxT), PAD, dtype=torch.long, device=DEVICE) | |
| for b, r in enumerate(rec_ids): | |
| lo, hi = int(off[r]), int(off[r+1]); T = hi - lo | |
| ptr = tok_ptr[lo:hi+1] | |
| nlo, nhi = int(ptr[0].item()), int(ptr[-1].item()) | |
| if nhi > nlo: | |
| counts = ptr[1:] - ptr[:-1] | |
| rows = torch.repeat_interleave(torch.arange(T, device=DEVICE), counts) | |
| x[b, rows, col[nlo:nhi]] = val[nlo:nhi].float() | |
| y[b, :T] = labels[lo:hi] | |
| return x, y | |
| def predict_rows(model, C, rec_ids, bs): | |
| model.eval(); rows = [] | |
| for s in range(0, len(rec_ids), bs): | |
| batch = rec_ids[s:s+bs] | |
| x, y = densify(C, batch) | |
| logits = model(x) | |
| pred = logits.argmax(-1) | |
| for b, r in enumerate(batch): | |
| T = int(C["offsets"][r+1]-C["offsets"][r]) | |
| rows.append({"labels": y[b, :T].cpu().numpy(), "predictions": pred[b, :T].cpu().numpy()}) | |
| return rows | |
| def train_arch(name, cfg, C): | |
| seed_everything(42, deterministic=True) | |
| mcfg = {"input_features": C["W"], "classes": len(TOPICS), **cfg["model"]} | |
| model = build_model(cfg["arch"], mcfg).to(DEVICE) | |
| opt = torch.optim.AdamW(model.parameters(), lr=cfg["lr"], weight_decay=0.01) | |
| loss_fn = nn.CrossEntropyLoss(ignore_index=PAD) | |
| train_ids = list(C["splits"]["train"]); val_ids = list(C["splits"]["validation"]); test_ids = list(C["splits"]["test"]) | |
| bs = cfg["bs"] | |
| best_f1, best_epoch, best_state = -1.0, 0, None | |
| params = sum(p.numel() for p in model.parameters()) | |
| print(f"\n=== {name} ({cfg['arch']}) params={params/1e6:.2f}M width={C['W']} ===", flush=True) | |
| for epoch in range(1, EPOCHS+1): | |
| model.train(); t0 = time.time() | |
| order = np.random.permutation(len(train_ids)) | |
| for st in range(0, len(order), bs): | |
| ids = [train_ids[i] for i in order[st:st+bs]] | |
| x, y = densify(C, ids) | |
| logits = model(x) | |
| loss = loss_fn(logits.reshape(-1, len(TOPICS)), y.reshape(-1)) | |
| opt.zero_grad(); loss.backward() | |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step() | |
| rows = predict_rows(model, C, val_ids, max(bs, 8)) | |
| truth, pred = flatten_labeled(rows) | |
| f1 = f1_score(truth, pred, average="macro", zero_division=0) | |
| print(f"{name} epoch {epoch:2d}/{EPOCHS}: val macro-F1={f1:.4f} ({time.time()-t0:.1f}s)", flush=True) | |
| if f1 > best_f1: | |
| best_f1, best_epoch = f1, epoch | |
| best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} | |
| model.load_state_dict(best_state) | |
| out = ROOT / cfg["out"]; out.mkdir(parents=True, exist_ok=True) | |
| torch.save({"state_dict": best_state, "architecture": cfg["arch"], "model": mcfg, | |
| "classes": TOPICS, "epoch": best_epoch, "selection": "best_validation_macro_f1"}, | |
| out / "checkpoint_best.pt") | |
| test_rows = predict_rows(model, C, test_ids, max(bs, 8)) | |
| m = metrics(test_rows, [5, 10]) | |
| (out / "results.json").write_text(json.dumps({"parameters": params, "best_epoch": best_epoch, | |
| "test": m}, indent=2) + "\n") | |
| print(f"{name} DONE best_epoch={best_epoch} test macroF1={m['macro_f1']:.4f} " | |
| f"IoU={m['topic_overlap']['macro_iou']:.4f}", flush=True) | |
| def main(): | |
| C = load_cache() | |
| for name, cfg in ARCHS.items(): | |
| train_arch(name, cfg, C) | |
| torch.cuda.empty_cache() | |
| print("\nALL DONE", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.08 kB
- Xet hash:
- 24b2e9a55797306b2f20021ab69d0fca3473ee326ccaa864e0c7e9b9d291585e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.