Buckets:
| """Losses for keypoint forecasters on the AIDE protocol (all in pixels). | |
| compute_losses(out, fut_xy, w, plain_loss) -> (total, {term: value}) | |
| out tensor [B,H,J,2] (plain model) -> only `plain_loss` (euc | l1 | mse) | |
| or dict {"pred":[B,H,J,2], "log_b":[B,H,J] (optional), "last":[B,1,J,2] (optional)} | |
| -> main term of kind `plain_loss` (default euc) + weighted nll / bone / jerk | |
| fut_xy [B,H,J,2] ground-truth future | |
| w weights {"euc","nll","bone","jerk"} (plan §4: 1 / 0.1 / 0.05 / 0.01); w["euc"] weights the | |
| main term whatever its kind ("pure L1" ablation = plain_loss="l1", w["nll"]=0) | |
| Terms | |
| euc mean Euclidean pixel error (== the metric) [or l1 / mse when plain_loss says so] | |
| nll isotropic Laplace NLL per joint: (|dx|+|dy|)/b + 2 log b, b = exp(log_b) px | |
| bone mean |‖p_i−p_j‖ − ‖g_i−g_j‖| over protocol.EDGES_BODY_HANDS (px) | |
| jerk mean ‖third difference‖ of the trajectory [last_obs, pred_1..H] (px) | |
| HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A), arms a1 / a2): compute_losses(..., jw=joint_weights(w_hand)) weights the | |
| 42 hand keypoints (PARTS lhand + rhand, HALPE 94-135) by w_hand in the Euclidean term only (other joints 1, not renormalised; NLL / | |
| bone / jerk and every metric unchanged); jw=None (the default, and joint_weights(1.0)) is the untouched code path. | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| DEFAULT_W = {"euc": 1.0, "nll": 0.1, "bone": 0.05, "jerk": 0.01} | |
| _EDGES = torch.as_tensor(P.EDGES_BODY_HANDS, dtype=torch.long) # [E,2] | |
| def euclidean(pred, gt, eps=1e-6, w=None): | |
| """w: optional weights broadcastable to [B,H,J] with mean 1 (horizon and/or per-clip weighting).""" | |
| e = torch.sqrt(((pred - gt) ** 2).sum(-1) + eps) | |
| return e.mean() if w is None else (e * w).mean() | |
| def joint_weights(hand_weight=1.0, device=None, dtype=torch.float32): | |
| """Pre-registered HM h=5 arms (a1) / (a2): per-joint weights [J] for the Euclidean term — the 42 hand keypoints (PARTS lhand + rhand, | |
| HALPE 94–135) get `hand_weight`, every other joint 1; NOT renormalised (the pre-registration fixes "other joints 1"). Returns None for | |
| hand_weight == 1 (or None) so that the default code path is untouched.""" | |
| if hand_weight is None or float(hand_weight) == 1.0: | |
| return None | |
| w = torch.ones(P.J, dtype=dtype, device=device) | |
| w[list(P.PARTS["lhand"]) + list(P.PARTS["rhand"])] = float(hand_weight) | |
| return w | |
| def plain(pred, gt, kind): | |
| if kind == "euc": | |
| return euclidean(pred, gt) | |
| if kind == "l1": | |
| return (pred - gt).abs().mean() | |
| if kind == "mse": | |
| return ((pred - gt) ** 2).mean() / 100.0 | |
| raise ValueError(kind) | |
| def laplace_nll(pred, gt, log_b, w=None): | |
| """Isotropic Laplace with scale b per (clip, horizon, joint). log_b: [B,H,J]. w as in euclidean().""" | |
| b = torch.exp(log_b) | |
| n = (pred - gt).abs().sum(-1) / b + 2.0 * log_b | |
| return n.mean() if w is None else (n * w).mean() | |
| def sample_weights(fut_xy, h_weights=None, motion_cap=0.0): | |
| """-> weights [B,H,1] with mean 1 (or None). h_weights: sequence of H horizon weights (normalised to | |
| mean 1). motion_cap > 0: per-clip weight = min(score / mean(score), cap), score = mean GT displacement | |
| over the future transitions and joints (the HM motion score), renormalised to mean 1 over the batch.""" | |
| B, H = fut_xy.shape[:2] | |
| w = fut_xy.new_ones(B, H, 1) | |
| used = False | |
| if h_weights is not None: | |
| hw = torch.as_tensor(list(h_weights), dtype=fut_xy.dtype, device=fut_xy.device) | |
| assert hw.numel() == H | |
| w = w * (hw / hw.mean())[None, :, None] | |
| used = True | |
| if motion_cap and motion_cap > 0: | |
| d = torch.sqrt(((fut_xy[:, 1:] - fut_xy[:, :-1]) ** 2).sum(-1) + 1e-6).mean((1, 2)) # [B] | |
| cw = torch.clamp(d / (d.mean() + 1e-6), max=motion_cap) | |
| cw = cw / cw.mean() | |
| w = w * cw[:, None, None] | |
| used = True | |
| return w if used else None | |
| def bone_length(pred, gt, edges=None): | |
| e = (_EDGES if edges is None else edges).to(pred.device) | |
| lp = torch.linalg.norm(pred[..., e[:, 0], :] - pred[..., e[:, 1], :], dim=-1) | |
| lg = torch.linalg.norm(gt[..., e[:, 0], :] - gt[..., e[:, 1], :], dim=-1) | |
| return (lp - lg).abs().mean() | |
| def jerk(pred, last=None): | |
| """Third temporal difference of [last, pred_1..H] (or of pred alone if last is None).""" | |
| traj = pred if last is None else torch.cat([last, pred], 1) | |
| if traj.shape[1] < 4: | |
| return pred.new_zeros(()) | |
| # successive differences (exactly zero for a constant trajectory in float32, unlike the | |
| # expanded 1/-3/3/-1 stencil which leaves ~1e-4 px of rounding at 1e3-px coordinates) | |
| d1 = traj[:, 1:] - traj[:, :-1] | |
| d2 = d1[:, 1:] - d1[:, :-1] | |
| d3 = d2[:, 1:] - d2[:, :-1] | |
| return torch.linalg.norm(d3, dim=-1).mean() | |
| def compute_losses(out, fut_xy, w=None, plain_loss="euc", sw=None, jw=None): | |
| """sw: optional sample weights [B,H,1] (mean 1) applied to the main Euclidean term and the NLL term. | |
| jw: optional per-joint weights [J] (joint_weights; pre-registered hand-joint weight) applied to the Euclidean main term ONLY — | |
| combined multiplicatively with sw when both are given; None (default) = unchanged code path.""" | |
| w = {**DEFAULT_W, **(w or {})} | |
| if jw is not None: | |
| assert plain_loss == "euc", "joint weights (--hand-weight) are defined for the Euclidean main term only" | |
| if torch.is_tensor(out): | |
| if jw is not None: | |
| l = euclidean(out, fut_xy, w=jw.to(device=out.device, dtype=out.dtype)[None, None, :]) | |
| else: | |
| l = plain(out, fut_xy, plain_loss) | |
| return l, {plain_loss: l.detach()} | |
| pred = out["pred"] | |
| if jw is not None: | |
| wt = jw.to(device=pred.device, dtype=pred.dtype)[None, None, :] # [1,1,J]; hands w_hand, others 1 | |
| terms = {"euc": euclidean(pred, fut_xy, w=wt if sw is None else sw * wt)} | |
| elif sw is not None and plain_loss == "euc": | |
| terms = {"euc": euclidean(pred, fut_xy, w=sw)} | |
| else: | |
| terms = {plain_loss: plain(pred, fut_xy, plain_loss)} # main term (euc by default) | |
| total = w["euc"] * terms[plain_loss] | |
| if "log_b" in out and out["log_b"] is not None and w["nll"] > 0: | |
| terms["nll"] = laplace_nll(pred, fut_xy, out["log_b"], w=sw) | |
| total = total + w["nll"] * terms["nll"] | |
| if w["bone"] > 0: | |
| terms["bone"] = bone_length(pred, fut_xy) | |
| total = total + w["bone"] * terms["bone"] | |
| if w["jerk"] > 0: | |
| terms["jerk"] = jerk(pred, out.get("last")) | |
| total = total + w["jerk"] * terms["jerk"] | |
| return total, {k: v.detach() for k, v in terms.items()} | |
| # ----------------------------------------------------------------------------- F7 ensemble distillation (pre-registered 2026-09-17) | |
| def distillation(out, tgt, w_coord=0.0, w_modes=0.0, w_logb=0.0, eps=1e-8): | |
| """Student output `out` (tensor or cwm3d dict) vs the mean-teacher target `tgt` = {"pred", "log_b" | None, "anchor_w" | None} | |
| (train_kp.teacher_target). Terms: coordinates = euclidean(pred_s, pred_t) (px, same reduction as the GT term); modes = KL(p_t || p_s) | |
| over the 3 anchor weights, mean over (B,H,J); log_b = L1(log_b_s, log_b_t). Every weight 0 -> (0, {}) without touching the outputs | |
| (the default code path). Terms whose inputs are missing (tensor-output models have no log_b / anchor_w) are skipped.""" | |
| if w_coord <= 0 and w_modes <= 0 and w_logb <= 0: | |
| pred = out["pred"] if isinstance(out, dict) else out | |
| return pred.new_zeros(()), {} | |
| pred = out["pred"] if isinstance(out, dict) else out | |
| total, terms = pred.new_zeros(()), {} | |
| if w_coord > 0: | |
| terms["distil"] = euclidean(pred, tgt["pred"]) | |
| total = total + w_coord * terms["distil"] | |
| if w_modes > 0 and isinstance(out, dict) and out.get("anchor_w") is not None and tgt.get("anchor_w") is not None: | |
| p_t, p_s = tgt["anchor_w"].clamp_min(eps), out["anchor_w"].clamp_min(eps) | |
| terms["distil_modes"] = (p_t * (torch.log(p_t) - torch.log(p_s))).sum(-1).mean() | |
| total = total + w_modes * terms["distil_modes"] | |
| if w_logb > 0 and isinstance(out, dict) and out.get("log_b") is not None and tgt.get("log_b") is not None: | |
| terms["distil_logb"] = (out["log_b"] - tgt["log_b"]).abs().mean() | |
| total = total + w_logb * terms["distil_logb"] | |
| return total, {k: v.detach() for k, v in terms.items()} | |
| # ----------------------------------------------------------------------------- K-v2 onset-hazard targets (plan A.1, 2026-09-14) | |
| # Per part p (P.PARTS order body/face/lhand/rhand) the centroid displacement d_h = ||centroid(fut_h) - centroid(last obs)|| (px); | |
| # "exceeded" at horizon h iff d_h > tau_p; onset = first exceedance, risk = not yet exceeded before h. The BCE below is the | |
| # on-branch of the BATON two-state chain (models/heads.py chain_bce with m = risk); P = 1 - prod(1 - lam) is forward_marginal with h_off = 0. | |
| _PART_IDX = [torch.as_tensor(P.PARTS[k], dtype=torch.long) for k in P.PART_NAMES] | |
| def hazard_stats(S10_train, q=0.75): | |
| """TRAIN clips on the protocol grid [N,10,J,3] (numpy) -> (tau [4] px, prior_P [4,H]) with H = T_PRED. | |
| tau_p = q-quantile over clips of the last-horizon displacement d_H; prior_P[p,h] = mean over clips of any_{k<=h}(d_k > tau_p).""" | |
| S = np.asarray(S10_train, dtype=np.float64) | |
| assert S.ndim == 4 and S.shape[1] == P.N_SAMPLE and S.shape[2] == P.J, S.shape | |
| last, fut = S[:, P.T_OBS - 1, :, :2], S[:, P.T_OBS:, :, :2] # [N,J,2], [N,H,J,2] | |
| tau = np.zeros(len(P.PART_NAMES)); prior = np.zeros((len(P.PART_NAMES), P.T_PRED)) | |
| for p, k in enumerate(P.PART_NAMES): | |
| idx = P.PARTS[k] | |
| d = np.linalg.norm(fut[:, :, idx].mean(2) - last[:, idx].mean(1)[:, None], axis=-1) # [N,H] | |
| tau[p] = float(np.quantile(d[:, -1], q)) | |
| prior[p] = np.maximum.accumulate(d > tau[p], axis=1).mean(0) | |
| return tau, prior | |
| def hazard_targets(last_xy, fut_xy, tau): | |
| """last_xy [B,J,2], fut_xy [B,H,J,2], tau [4] -> (onset [B,4,H], risk [B,4,H]) float tensors. | |
| ex = d > tau; cum = cummax(ex); prev = cum shifted right by one (0 at h = 1); onset = ex * (1 - prev); risk = 1 - prev.""" | |
| tau = torch.as_tensor(tau, dtype=fut_xy.dtype, device=fut_xy.device) | |
| ex = [] | |
| for p, idx in enumerate(_PART_IDX): | |
| idx = idx.to(fut_xy.device) | |
| d = torch.linalg.norm(fut_xy[:, :, idx].mean(2) - last_xy[:, idx].mean(1)[:, None], dim=-1) # [B,H] | |
| ex.append(d > tau[p]) | |
| ex = torch.stack(ex, 1).to(fut_xy.dtype) # [B,4,H] | |
| cum = torch.cummax(ex, dim=-1).values | |
| prev = torch.cat([torch.zeros_like(cum[..., :1]), cum[..., :-1]], -1) | |
| return ex * (1 - prev), 1 - prev | |
| def hazard_bce(logits, onset, risk): | |
| """Risk-weighted Bernoulli NLL of the onset hazards: sum(BCE(logits, onset) * risk) / max(sum(risk), 1).""" | |
| l = torch.nn.functional.binary_cross_entropy_with_logits(logits.float(), onset.float(), reduction="none") | |
| r = risk.float() | |
| return (l * r).sum() / r.sum().clamp_min(1.0) | |
Xet Storage Details
- Size:
- 11.3 kB
- Xet hash:
- 08497e54ed43eb4067cdfd91a3f56c1df6bebbb947ea69dfee6ec6ddb08bc366
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.