Buckets:
| """cwm3/probes/vjepa2_temporal_extract.py — P1b: re-extract V-JEPA2 (facebook/vjepa2-vitg-fpc64-384) for a DIAGNOSTIC subset of <= 2,000 | |
| front-camera windows, keeping the within-window TIME dimension that the frontfeat cache averages away. | |
| Windows are the cache's own windows (the last bank slot of a probe anchor: END time tau = floor(2t)/2 on the frame/20 axis, frames at | |
| round((tau + rel) * 20), rel = linspace(-2, 0, 16)); decoding / resizing / normalisation are copied VERBATIM from s12_wm/extract_crop.py | |
| (--cam front: decord decode at 384x384, /255, resize 438 antialias, center-crop 384, ImageNet mean/std, bf16), so the time-mean of the new | |
| tensor must reproduce the stored `region` tokens — asserted per window (cos) and reported. | |
| Stored per window: tt f16[8, 16, 1408] = 8 tubelet time groups x 4x4 spatial (6x6 patch pooling of the 24x24 grid) — i.e. the [n, 8, 24, 24, 1408] | |
| hidden states before the `.mean(1)` of extract_crop.py; g f16[8, 1408] = per-time-group spatial mean. | |
| Selection (seeded): balanced over regime x split, device-capped, half of the windows enriched with road pressure (P0 active or a P onset within | |
| 3 s), only routes whose qcamera.mp4 is local with the canonical size (results/baton_leaves.json). Test-split windows are processed locally only. | |
| GPU used directly (coordinator rule: no lock; free VRAM >= need + 4 GB checked), threads <= 4. Output: results/probes/vjepa2_temporal/ | |
| windows.npz {tt, g, cos_check, uid, t_row, t_anchor, regime, split, row, dongle, ...}, run json (GPU minutes, bytes, model id). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| import common as C # noqa: E402 | |
| OUT = C.OUT / "vjepa2_temporal" | |
| BATON_ROOT = Path("/home/USER/BATON") | |
| ENCODER_HF = "facebook/vjepa2-vitg-fpc64-384" | |
| CLIP_FRAMES, CLIP_SPAN = 16, 2.0 | |
| SIZE, RESIZE_TO = 384, 438 | |
| BATCH = 8 | |
| LOG = OUT / "extract.log" | |
| def log(m): | |
| C.log(m, LOG) | |
| def uid_to_leaf(uid): | |
| model, dongle, route_id, leaf = uid.split("__") | |
| return BATON_ROOT / model / dongle / route_id / leaf | |
| # ----------------------------------------------------------------------------- selection | |
| def select_windows(n_total, seed=0, cap_per_dongle=20): | |
| leaves = json.loads((C.CWM3 / "results" / "baton_leaves.json").read_text())["leaves"] | |
| with np.load(C.BANKS / "front_tok.npz", allow_pickle=True) as z: | |
| bank_t, names, off = z["t"].astype(np.float64), [str(u) for u in z["uid_index"]], z["route_off"] | |
| rng = np.random.default_rng(seed) | |
| quota = {"train": 0.5, "val": 0.25, "test": 0.25} | |
| picked = [] | |
| for regime in C.REGIMES: | |
| for split in C.SPLITS: | |
| raw = C.load_raw(regime, split) | |
| T = C.build_targets(raw) | |
| n = len(raw["t"]) | |
| ti = raw["tok_idx"][:, -1].astype(np.int64) | |
| ok = ti >= 0 | |
| # local canonical qcamera | |
| leaf_key = np.array(["/".join(u.split("__")) for u in raw["uid"]]) | |
| local = np.array([leaves.get(k, {}).get("local_status") == "local_identical_size" and (uid_to_leaf(u) / "qcamera.mp4").exists() | |
| for k, u in zip(leaf_key, raw["uid"])]) | |
| ok &= local | |
| pos = np.zeros(n, bool) | |
| for k, v in T.items(): | |
| if k.startswith("cur_P0_") or k.startswith("fut_P_"): | |
| pos |= v["y"] & v["m"] | |
| want = int(round(n_total / 2 * quota[split])) | |
| for enriched, share in ((True, 0.5), (False, 0.5)): | |
| cand = np.flatnonzero(ok & (pos if enriched else ~pos)) | |
| rng.shuffle(cand) | |
| per_d = {} | |
| got = 0 | |
| for i in cand: | |
| d = raw["dongle"][i] | |
| if per_d.get(d, 0) >= cap_per_dongle: | |
| continue | |
| per_d[d] = per_d.get(d, 0) + 1 | |
| picked.append({"regime": regime, "split": split, "row": int(i), "uid": str(raw["uid"][i]), "dongle": str(d), "t_anchor": float(raw["t"][i]), | |
| "t_row": float(bank_t[ti[i]]), "bank_row": int(ti[i]), "enriched": bool(enriched)}) | |
| got += 1 | |
| if got >= int(round(want * share)): | |
| break | |
| log(f"[select] {regime}/{split}: candidates local={local.mean():.3f} picked so far {len(picked)}") | |
| # de-duplicate identical windows (same uid, t_row) | |
| seen, uniq = set(), [] | |
| for p in picked: | |
| k = (p["uid"], round(p["t_row"], 2)) | |
| if k not in seen: | |
| seen.add(k) | |
| uniq.append(p) | |
| return uniq[:n_total] | |
| # ----------------------------------------------------------------------------- extraction | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--n", type=int, default=2000) | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--limit-routes", type=int, default=0) | |
| ap.add_argument("--select-only", action="store_true") | |
| a = ap.parse_args() | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| t0 = time.time() | |
| sel_p = OUT / "selection.json" | |
| if sel_p.exists(): | |
| sel = json.loads(sel_p.read_text()) | |
| else: | |
| sel = select_windows(a.n, a.seed) | |
| sel_p.write_text(json.dumps(sel)) | |
| log(f"{len(sel)} windows selected over {len(set(p['uid'] for p in sel))} routes / {len(set(p['dongle'] for p in sel))} dongles; " | |
| + ", ".join(f"{r}/{s}={sum(1 for p in sel if p['regime'] == r and p['split'] == s)}" for r in C.REGIMES for s in C.SPLITS) | |
| + f"; enriched={sum(p['enriched'] for p in sel)}") | |
| if a.select_only: | |
| return | |
| import torch | |
| import decord | |
| from torchvision.transforms.v2 import functional as TVF | |
| from transformers import VJEPA2Model | |
| free, total = torch.cuda.mem_get_info() | |
| assert free / 2**30 >= 8 + 4, f"free VRAM {free / 2**30:.1f} GB < need 8 + 4 GB" | |
| torch.set_num_threads(4) | |
| model = VJEPA2Model.from_pretrained(ENCODER_HF, torch_dtype=torch.bfloat16, attn_implementation="sdpa").cuda().eval() | |
| MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda() | |
| STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda() | |
| decord.bridge.set_bridge("native") | |
| rel = np.linspace(-CLIP_SPAN, 0.0, CLIP_FRAMES) | |
| N = len(sel) | |
| part = OUT / "partial.npz" | |
| tt_path = OUT / "tt.npy" | |
| if part.exists() and tt_path.exists(): # resume (the host killed an earlier pass for memory) | |
| z = np.load(part) | |
| g, cos_check, done, gpu_s = z["g"], z["cos_check"], z["done"].astype(bool), float(z["gpu_s"]) | |
| tt = np.lib.format.open_memmap(tt_path, mode="r+") | |
| assert tt.shape == (N, 8, 16, 1408) | |
| log(f"resume: {done.sum()} / {N} windows already done (gpu {gpu_s / 60:.1f} min)") | |
| else: | |
| tt = np.lib.format.open_memmap(tt_path, mode="w+", dtype=np.float16, shape=(N, 8, 16, 1408)) | |
| g = np.zeros((N, 8, 1408), np.float16) | |
| cos_check = np.full(N, np.nan, np.float32) | |
| done = np.zeros(N, bool) | |
| gpu_s = 0.0 | |
| by_route = {} | |
| for j, p in enumerate(sel): | |
| by_route.setdefault(p["uid"], []).append(j) | |
| routes = sorted(by_route) | |
| if a.limit_routes: | |
| routes = routes[:a.limit_routes] | |
| n_clips = 0 | |
| def save_partial(): | |
| tt.flush() | |
| np.savez(part, g=g, cos_check=cos_check, done=done, gpu_s=np.array(gpu_s)) | |
| with torch.no_grad(): | |
| for ri, uid in enumerate(routes): | |
| js = [j for j in by_route[uid] if not done[j]] | |
| if not js: | |
| continue | |
| try: | |
| vr = decord.VideoReader(str(uid_to_leaf(uid) / "qcamera.mp4"), ctx=decord.cpu(0), width=SIZE, height=SIZE, num_threads=4) | |
| except Exception as e: | |
| log(f" fail open {uid[:50]}: {str(e)[:80]}") | |
| continue | |
| nfr = len(vr) | |
| idx_all = [np.clip(np.round((sel[j]["t_row"] + rel) * 20).astype(int), 0, nfr - 1) for j in js] | |
| # stored region tokens for the check | |
| with np.load(C.FF / f"{uid}.npz") as z: | |
| t_ff = z["t"].astype(np.float64) | |
| pos = {round(float(t), 2): k for k, t in enumerate(t_ff)} | |
| want = [pos.get(round(sel[j]["t_row"], 2)) for j in js] | |
| reg_ff = {k: z["region"][k].astype(np.float32) for k in set(w for w in want if w is not None)} | |
| for b0 in range(0, len(js), BATCH): | |
| jb = js[b0:b0 + BATCH] | |
| idxs = idx_all[b0:b0 + BATCH] | |
| uniq = np.unique(np.concatenate(idxs)) | |
| frames = vr.get_batch(list(uniq)).asnumpy() # [F,384,384,3] RGB (full frame, x0 = 0) | |
| posf = {int(u): k for k, u in enumerate(uniq)} | |
| arr = np.stack([frames[[posf[int(v)] for v in ii]] for ii in idxs]) # [n,16,384,384,3] | |
| x = torch.from_numpy(arr).cuda(non_blocking=True) | |
| n = x.shape[0] | |
| x = x.permute(0, 1, 4, 2, 3).reshape(n * CLIP_FRAMES, 3, SIZE, SIZE).float().div_(255.0) | |
| x = TVF.resize(x, [RESIZE_TO], antialias=True) | |
| x = TVF.center_crop(x, [SIZE, SIZE]) | |
| x = ((x - MEAN) / STD).view(n, CLIP_FRAMES, 3, SIZE, SIZE).to(torch.bfloat16) | |
| torch.cuda.synchronize() | |
| tg = time.time() | |
| out = model(pixel_values_videos=x, skip_predictor=True) | |
| h = out.last_hidden_state.float() # [n,4608,1408] = 8 x 24 x 24 | |
| torch.cuda.synchronize() | |
| gpu_s += time.time() - tg | |
| r = h.view(n, 8, 24, 24, 1408).view(n, 8, 4, 6, 4, 6, 1408).mean((3, 5)) # [n,8,4,4,1408] | |
| r = r.reshape(n, 8, 16, 1408) | |
| tt[jb] = r.cpu().numpy().astype(np.float16) | |
| g[jb] = r.mean(2).cpu().numpy().astype(np.float16) | |
| rt = r.mean(1).cpu().numpy() # time-mean -> must equal the stored region | |
| for k, j in enumerate(jb): | |
| w = want[b0 + k] | |
| if w is not None: | |
| a_, b_ = rt[k].reshape(-1), reg_ff[w].reshape(-1) | |
| cos_check[j] = float((a_ * b_).sum() / (np.linalg.norm(a_) * np.linalg.norm(b_) + 1e-6)) | |
| done[j] = True | |
| n_clips += n | |
| del vr | |
| if ri % 20 == 0: | |
| save_partial() | |
| log(f"[extract] {ri + 1}/{len(routes)} routes, {int(done.sum())} windows done, gpu {gpu_s / 60:.1f} min, cos_check median {np.nanmedian(cos_check[done]):.4f} ({time.time() - t0:.0f}s)") | |
| save_partial() | |
| meta = {k: np.array([p[k] for p in sel]) for k in ("regime", "split", "row", "uid", "dongle", "t_anchor", "t_row", "bank_row", "enriched")} | |
| np.savez(OUT / "windows.npz", g=g, cos_check=cos_check, done=done, tt_file=np.array(str(OUT / "tt.npy")), **meta, | |
| model=np.array(ENCODER_HF), layout=np.array("tt[N,8,16,1408]: 8 tubelet time groups x (4x4 spatial, 6x6-patch pooled) x 1408; g = spatial mean per group")) | |
| rec = {"n_selected": N, "n_done": int(done.sum()), "routes": len(routes), "gpu_minutes_forward": round(gpu_s / 60, 2), "wall_minutes": round((time.time() - t0) / 60, 1), | |
| "bytes_tt": int(N * 8 * 16 * 1408 * 2), "bytes_written_total": int(sum(p.stat().st_size for p in OUT.glob("*.np*"))), | |
| "peak_vram_GB": round(torch.cuda.max_memory_allocated() / 2**30, 2), "cos_check_median": float(np.nanmedian(cos_check[done])), | |
| "cos_check_p01": float(np.nanpercentile(cos_check[done], 1)), "cos_check_n": int(np.isfinite(cos_check).sum()), | |
| "model": ENCODER_HF, "created_utc": C.stamp(), "test_windows_processed_locally": int((meta["split"] == "test").sum())} | |
| C.write_json(OUT / f"run_extract_{rec['created_utc']}.json", rec) | |
| C.write_json(OUT / "run_extract_latest.json", rec) | |
| log(f"DONE {rec}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12.2 kB
- Xet hash:
- 0a9cdcd08052088b3340dc1cb312802e14703f042b3dfe665e393fe1b4f0c630
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.