From 22f603d81f9868148ef5af3f3dc3dd810fe9814a Mon Sep 17 00:00:00 2001 From: ruv Date: Thu, 11 Jun 2026 00:41:05 -0400 Subject: [PATCH] =?UTF-8?q?feat(benchmarks):=20efficiency=20sweep=20MEASUR?= =?UTF-8?q?ED=20=E2=80=94=20half=20model=20dominates=20full=20reference?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compact WiFlow-STD variants on the same data/split/protocol: half (843,834 params, 0.38x) strictly dominates the 2.23M reference (PCK@20 96.62 vs 96.61, PCK@50 99.47 vs 99.11, MPJPE 0.00898 vs 0.0094) — the published architecture is over-parameterized for its own benchmark. quarter (338k) 96.05%; tiny (56,290 params, 1/39.5) holds 94.11% — a ~220KB fp32 edge candidate. In-domain caveats recorded; cross-domain untested. Co-Authored-By: claude-flow --- benchmarks/wiflow-std/RESULTS.md | 34 ++ .../wiflow-std/remote/sweep/model_compact.py | 332 ++++++++++++++++++ .../wiflow-std/remote/sweep/run_sweep.py | 259 ++++++++++++++ .../wiflow-std/results/efficiency_sweep.jsonl | 3 + 4 files changed, 628 insertions(+) create mode 100644 benchmarks/wiflow-std/remote/sweep/model_compact.py create mode 100644 benchmarks/wiflow-std/remote/sweep/run_sweep.py create mode 100644 benchmarks/wiflow-std/results/efficiency_sweep.jsonl diff --git a/benchmarks/wiflow-std/RESULTS.md b/benchmarks/wiflow-std/RESULTS.md index 673df59e..a761216c 100644 --- a/benchmarks/wiflow-std/RESULTS.md +++ b/benchmarks/wiflow-std/RESULTS.md @@ -233,6 +233,40 @@ for size → static QDQ conv-only (Percentile or MinMax, `results/retrained_int8_static_percentile_conv.onnx`), which strictly dominates dynamic int8 on accuracy at ~equal latency and +0.09 MB. +## Efficiency sweep (MEASURED, overnight 2026-06-10/11) + +ADR-152 beyond-SOTA track: compact purpose-built variants of the WiFlow-STD +architecture, trained from scratch on the same cleaned dataset, identical +seed-42 file-level split, loss and protocol as the measurement-(a) reference +(fp32, batch 64, ≤50 epochs, patience 5; RTX 5080, ~22–29 min/variant). +Variant transforms are pure channel/group/stride scalings of an +architecture-exact parameterized model (validated: reproduces 2,225,042 params +at the reference config). Scripts: `remote/sweep/`; raw: +`results/efficiency_sweep.jsonl`; checkpoints `results/{half,quarter,tiny}_best.pth` +(gitignored). + +| Variant | Params | vs 2.23M | Clean-test PCK@20 | PCK@50 | MPJPE | Best epoch | +|---|---|---|---|---|---|---| +| full (reference, meas. a) | 2,225,042 | 1× | 96.61% | 99.11% | 0.0094 | 36 | +| **half** | **843,834** | **0.38×** | **96.62%** | **99.47%** | **0.00898** | 23 | +| quarter | 338,600 | 0.15× | 96.05% | 99.43% | 0.00928 | 50 | +| tiny | 56,290 | 0.025× | 94.11% | 99.36% | 0.0125 | 47 | + +Findings: + +- **The half model (843k params) strictly dominates the full reference** on + this dataset — equal PCK@20, better PCK@50 and MPJPE, converges in fewer + epochs. The published 2.23M architecture is over-parameterized for its own + benchmark. +- **tiny (56k params, 1/39.5) holds 94.11% PCK@20** — a ~220 KB fp32 / + ~60 KB int8-class model in reach of severely constrained edge targets, + at −2.5 pt from the full reference. +- Caveats: in-domain (5-subject random-file split) like every number on this + dataset; single run per variant; corruption-free test subset (52,560). + Cross-domain behavior of compact variants is untested — ADR-150's evidence + says capacity *hurts* cross-subject, so the compact end may generalize no + worse, but that is a hypothesis, not a measurement. + ## Measurement (b): BLOCKED-ON-DATA (attempted 2026-06-10) The fine-tune-on-ESP32 measurement stopped at dataset characterization, per the diff --git a/benchmarks/wiflow-std/remote/sweep/model_compact.py b/benchmarks/wiflow-std/remote/sweep/model_compact.py new file mode 100644 index 00000000..b571de37 --- /dev/null +++ b/benchmarks/wiflow-std/remote/sweep/model_compact.py @@ -0,0 +1,332 @@ +"""Configurable compact variants of the WiFlow-STD pose model (ADR-152 efficiency sweep). + +This is a parameterized copy of upstream models/{pose_model,tcn,convnet,attention}.py +(DY2434/WiFlow @ 06899d29, Apache-2.0). upstream/ is NOT modified. Deviations from +upstream, all forced by shrinking channels and documented per variant in run_sweep.py: + +1. TCN grouped-conv groups: upstream hardcodes groups=20, which does not divide + the compact channel counts (e.g. 270, 135, 85). Rule here: + - groups_mode='gcd20': per-conv groups = gcd(channels, 20) (== 20 wherever + upstream's choice is valid, incl. the 540-ch input conv; falls back to the + largest common divisor with 20 otherwise). + - groups_mode='depthwise': groups = channels (tiny variant only). +2. Conv2d downsampling strides: upstream uses 4 stride-(1,2) blocks because + 240/2^4 = 15 == n_keypoints. With smaller TCN output widths that would leave + <15 rows and AdaptiveAvgPool2d((15,1)) would duplicate rows across keypoints. + Rule: halve the width only while the result stays >= 15 (stride-2 blocks + first, stride-1 after). Full model: 240 -> 4 halvings = upstream exactly. +3. input_pw_groups (tiny only): the dense 540->c pointwise + residual downsample + in TCN block 1 cost 2*540*c params (a ~117k floor that alone exceeds the + tiny <100k budget). tiny groups these two convs (groups=4; 4 | gcd(540, 68)). +4. Decoder mid-channels: upstream 64->32; here c_last -> max(c_last // 2, 4). +""" +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def tcn_groups(channels: int, mode: str) -> int: + if mode == 'depthwise': + return channels + if mode == 'gcd20': + return math.gcd(channels, 20) + raise ValueError(mode) + + +# ---------------------------------------------------------------- TCN (copy of tcn.py) +class Chomp1d(nn.Module): + def __init__(self, chomp_size): + super().__init__() + self.chomp_size = chomp_size + + def forward(self, x): + return x[:, :, :-self.chomp_size].contiguous() + + +class CompactGroupedTemporalBlock(nn.Module): + """Upstream InnerGroupedTemporalBlock with parameterized groups.""" + + def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, + dropout=0.2, groups_mode='gcd20', pw_groups=1): + super().__init__() + g_in = tcn_groups(n_inputs, groups_mode) + g_out = tcn_groups(n_outputs, groups_mode) + self.groups = (g_in, g_out) + self.pw_groups = pw_groups + + self.conv1_group = nn.Conv1d(n_inputs, n_inputs, kernel_size, stride=stride, + padding=padding, dilation=dilation, + groups=g_in, bias=False) + self.chomp1 = Chomp1d(padding) if padding > 0 else nn.Identity() + self.bn1_group = nn.BatchNorm1d(n_inputs) + self.relu1_group = nn.SiLU(inplace=True) + + self.conv1_pw = nn.Conv1d(n_inputs, n_outputs, 1, groups=pw_groups, bias=False) + self.bn1_pw = nn.BatchNorm1d(n_outputs) + self.relu1_pw = nn.SiLU(inplace=True) + self.dropout1 = nn.Dropout(dropout) + + self.conv2_group = nn.Conv1d(n_outputs, n_outputs, kernel_size, stride=1, + padding=padding, dilation=dilation, + groups=g_out, bias=False) + self.chomp2 = Chomp1d(padding) if padding > 0 else nn.Identity() + self.bn2_group = nn.BatchNorm1d(n_outputs) + self.relu2_group = nn.SiLU(inplace=True) + + self.conv2_pw = nn.Conv1d(n_outputs, n_outputs, 1, bias=False) + self.bn2_pw = nn.BatchNorm1d(n_outputs) + self.relu2_pw = nn.SiLU(inplace=True) + self.dropout2 = nn.Dropout(dropout) + + self.downsample = nn.Sequential( + nn.Conv1d(n_inputs, n_outputs, 1, groups=pw_groups, bias=False), + nn.BatchNorm1d(n_outputs) + ) if n_inputs != n_outputs else nn.Identity() + + def forward(self, x): + res = self.downsample(x) + out = self.conv1_group(x) + out = self.chomp1(out) + out = self.bn1_group(out) + out = self.relu1_group(out) + out = self.conv1_pw(out) + out = self.bn1_pw(out) + out = self.relu1_pw(out) + out = self.dropout1(out) + out = self.conv2_group(out) + out = self.chomp2(out) + out = self.bn2_group(out) + out = self.relu2_group(out) + out = self.conv2_pw(out) + out = self.bn2_pw(out) + out = self.relu2_pw(out) + out = self.dropout2(out) + return F.silu(out + res) + + +class CompactTemporalBlock(nn.Module): + def __init__(self, num_inputs, num_channels, kernel_size=3, dropout=0.2, + groups_mode='gcd20', input_pw_groups=1): + super().__init__() + layers = [] + for i, out_channels in enumerate(num_channels): + dilation_size = 2 ** i + in_channels = num_inputs if i == 0 else num_channels[i - 1] + layers.append(CompactGroupedTemporalBlock( + in_channels, out_channels, kernel_size, stride=1, + dilation=dilation_size, padding=(kernel_size - 1) * dilation_size, + dropout=dropout, groups_mode=groups_mode, + pw_groups=input_pw_groups if i == 0 else 1)) + self.network = nn.Sequential(*layers) + + def forward(self, x): + return self.network(x) + + +# ------------------------------------------------------- Conv2d path (copy of convnet.py) +class AsymmetricConvBlock(nn.Module): + """Upstream block with parameterized width stride (upstream: always (1,2)).""" + + def __init__(self, in_channels, out_channels, dropout=0.3, stride_w=2): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=(1, 3), + stride=(1, stride_w), padding=(0, 1)), + nn.BatchNorm2d(out_channels), + nn.SiLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)), + nn.BatchNorm2d(out_channels), + nn.SiLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)), + nn.BatchNorm2d(out_channels) + ) + self.downsample = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, + stride=(1, stride_w), bias=False), + nn.BatchNorm2d(out_channels) + ) + self.activation = nn.SiLU(inplace=True) + + def forward(self, x): + return self.activation(self.block(x) + self.downsample(x)) + + +class ConvBlock1(nn.Module): + def __init__(self, in_channels, out_channels, dropout=0.3): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)), + nn.BatchNorm2d(out_channels), + nn.SiLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)), + nn.BatchNorm2d(out_channels), + nn.SiLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)), + nn.BatchNorm2d(out_channels) + ) + self.downsample = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + self.activation = nn.SiLU(inplace=True) + + def forward(self, x): + return self.activation(self.block(x) + self.downsample(x)) + + +# ----------------------------------------------------- attention (verbatim attention.py) +class AxialAttention(nn.Module): + def __init__(self, in_planes, out_planes, groups=8, stride=1, bias=False, width=False): + assert (in_planes % groups == 0) and (out_planes % groups == 0) + super().__init__() + self.in_planes = in_planes + self.out_planes = out_planes + self.groups = groups + self.group_planes = out_planes // groups + self.stride = stride + self.bias = bias + self.width = width + self.qkv_transform = nn.Conv1d(in_planes, out_planes * 3, kernel_size=1, + stride=1, padding=0, bias=False) + self.bn_qkv = nn.BatchNorm1d(out_planes * 3) + self.bn_similarity = nn.BatchNorm2d(groups) + self.bn_output = nn.BatchNorm1d(out_planes) + if stride > 1: + self.pooling = nn.AvgPool2d(stride, stride=stride) + nn.init.normal_(self.qkv_transform.weight.data, 0, math.sqrt(1. / self.in_planes)) + + def forward(self, x): + if self.width: + x = x.permute(0, 2, 1, 3) + else: + x = x.permute(0, 3, 1, 2) + N, W, C, H = x.shape + x = x.contiguous().view(N * W, C, H) + qkv = self.bn_qkv(self.qkv_transform(x)) + qkv = qkv.reshape(N * W, 3, self.out_planes, H).permute(1, 0, 2, 3) + q, k, v = qkv[0], qkv[1], qkv[2] + q = q.reshape(N * W, self.groups, self.group_planes, H) + k = k.reshape(N * W, self.groups, self.group_planes, H) + v = v.reshape(N * W, self.groups, self.group_planes, H) + qk = torch.einsum('bgci, bgcj->bgij', q, k) + qk = self.bn_similarity(qk) + similarity = F.softmax(qk, dim=-1) + sv = torch.einsum('bgij,bgcj->bgci', similarity, v) + sv = sv.reshape(N * W, self.out_planes, H) + out = self.bn_output(sv) + out = out.view(N, W, self.out_planes, H) + if self.width: + out = out.permute(0, 2, 1, 3) + else: + out = out.permute(0, 2, 3, 1) + if self.stride > 1: + out = self.pooling(out) + return out + + +class DualAxialAttention(nn.Module): + def __init__(self, in_planes, out_planes, groups=8, stride=1, bias=False): + super().__init__() + self.width_axis = AxialAttention(in_planes, out_planes, groups, stride, bias, width=True) + self.height_axis = AxialAttention(out_planes, out_planes, groups, stride, bias, width=False) + + def forward(self, x): + return self.height_axis(self.width_axis(x)) + + +# --------------------------------------------------------------- full model +def compute_strides(width: int, n_blocks: int, target: int = 15): + """Halve width while result stays >= target (upstream: 240 -> 4 halvings -> 15).""" + strides = [] + for _ in range(n_blocks): + nxt = (width + 1) // 2 # conv k=3 s=2 p=1: out = ceil(in/2) + if nxt >= target: + strides.append(2) + width = nxt + else: + strides.append(1) + return strides, width + + +class CompactWiFlowPoseModel(nn.Module): + """Parameterized upstream WiFlowPoseModel. + + Upstream config == tcn_channels=[540,440,340,240], conv_channels=[8,16,32,64], + attn_groups=8, groups_mode='gcd20' (gcd(c,20)==20 for all upstream channels), + input_pw_groups=1 -> identical architecture, 2,225,042 params. + """ + + def __init__(self, tcn_channels, conv_channels, attn_groups, + groups_mode='gcd20', input_pw_groups=1, dropout=0.3, + num_subcarriers=540, num_keypoints=15): + super().__init__() + self.tcn = CompactTemporalBlock( + num_inputs=num_subcarriers, num_channels=tcn_channels, kernel_size=3, + dropout=dropout, groups_mode=groups_mode, input_pw_groups=input_pw_groups) + + self.up = ConvBlock1(1, conv_channels[0]) + + strides, self.final_width = compute_strides( + tcn_channels[-1], len(conv_channels), target=num_keypoints) + self.conv_strides = strides + self.residual_blocks = nn.ModuleList() + in_channels = conv_channels[0] + for out_channels, s in zip(conv_channels, strides): + self.residual_blocks.append( + AsymmetricConvBlock(in_channels, out_channels, stride_w=s)) + in_channels = out_channels + + c_last = conv_channels[-1] + self.attention = DualAxialAttention(c_last, c_last, groups=attn_groups) + + c_mid = max(c_last // 2, 4) + self.decoder = nn.Sequential( + nn.Conv2d(c_last, c_mid, kernel_size=3, padding=1), + nn.BatchNorm2d(c_mid), + nn.SiLU(inplace=True), + nn.Conv2d(c_mid, 2, kernel_size=1), + nn.BatchNorm2d(2), + nn.SiLU(inplace=True) + ) + self.avg_pool = nn.AdaptiveAvgPool2d((num_keypoints, 1)) + self._initialize_weights() + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, (nn.BatchNorm1d, nn.LayerNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.xavier_normal_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + # [B, 540, 20] + x = self.tcn(x) # [B, C_tcn, 20] + x = x.transpose(1, 2).unsqueeze(1) # [B, 1, 20, C_tcn] + x = self.up(x) + for block in self.residual_blocks: + x = block(x) # [B, C_conv, 20, W'] + x = x.permute(0, 1, 3, 2) # [B, C_conv, W', 20] + x = self.attention(x) + x = self.decoder(x) # [B, 2, W', 20] + x = self.avg_pool(x).squeeze(-1) # [B, 2, 15] + return x.transpose(1, 2) # [B, 15, 2] + + +def describe(model: 'CompactWiFlowPoseModel'): + params = sum(p.numel() for p in model.parameters()) + tcn_g = [blk.groups for blk in model.tcn.network] + return {'params': params, 'tcn_groups_per_block': tcn_g, + 'conv_strides': model.conv_strides, 'final_width': model.final_width} diff --git a/benchmarks/wiflow-std/remote/sweep/run_sweep.py b/benchmarks/wiflow-std/remote/sweep/run_sweep.py new file mode 100644 index 00000000..b0312d31 --- /dev/null +++ b/benchmarks/wiflow-std/remote/sweep/run_sweep.py @@ -0,0 +1,259 @@ +"""WiFlow-STD compact-variant efficiency sweep (ADR-152) — sequential overnight runner. + +Trains compact variants of the upstream WiFlow-STD architecture on the same +data/split as the full-size reference retraining (seed 42, file-level 70/15/15, +upstream dataset.py) and evaluates PCK@10..50 + MPJPE on the full test split and +the corruption-free test subset (file indices < 487). + +Training mirrors upstream run.py/train.py defaults except: +- fp32 only (no fp16 autocast / GradScaler — avoids the BN-poisoning trap + documented in RESULTS.md defect 5; data on disk is already cleaned). +- batch 64 (kept modest: another GPU job may share the 16 GB card tonight). +- scheduler + early stopping keyed on val MPJPE (upstream early-stops on val MPE + with patience 5; same here). + +Usage: + venv/bin/python sweep/run_sweep.py --dry-run # param counts only + nohup venv/bin/python sweep/run_sweep.py > sweep/sweep.log 2>&1 & + +Idempotent: variants already present in sweep/results.jsonl are skipped. +""" +import argparse +import copy +import json +import os +import random +import sys +import time + +import numpy as np +import torch +from torch.utils.data import DataLoader, Subset + +BENCH = os.path.expanduser('~/wiflow-std-bench') +SWEEP = os.path.join(BENCH, 'sweep') +sys.path.insert(0, os.path.join(BENCH, 'upstream')) +sys.path.insert(0, SWEEP) + +from dataset import PreprocessedCSIKeypointsDataset, create_preprocessed_train_val_test_loaders # noqa: E402 +from losses.pose_loss import PoseLoss # noqa: E402 +from utils.metrics import calculate_pck, calculate_mpjpe # noqa: E402 +from model_compact import CompactWiFlowPoseModel, describe # noqa: E402 + +VARIANTS = [ + # name, tcn_channels, conv_channels, attn_groups, groups_mode, input_pw_groups + dict(name='half', tcn=[270, 220, 170, 120], conv=[4, 8, 16, 32], attn_groups=4, + groups_mode='gcd20', input_pw_groups=1), + dict(name='quarter', tcn=[135, 110, 85, 60], conv=[2, 4, 8, 16], attn_groups=2, + groups_mode='gcd20', input_pw_groups=1), + dict(name='tiny', tcn=[68, 56, 44, 32], conv=[2, 4, 8, 16], attn_groups=2, + groups_mode='depthwise', input_pw_groups=4), +] + +BATCH = 64 +EPOCHS = 50 +PATIENCE = 5 +LR = 1e-4 +WEIGHT_DECAY = 5e-5 +SEED = 42 +CORRUPT_FILE_START = 487 # files 487-499 were zero-filled by clean_nan.py + + +def set_seed(seed=SEED): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def build_model(v, dropout=0.5): + return CompactWiFlowPoseModel( + tcn_channels=v['tcn'], conv_channels=v['conv'], attn_groups=v['attn_groups'], + groups_mode=v['groups_mode'], input_pw_groups=v['input_pw_groups'], + dropout=dropout) + + +@torch.no_grad() +def evaluate(model, loader, device): + model.eval() + totals = {t: 0.0 for t in (0.1, 0.2, 0.3, 0.4, 0.5)} + total_mpe, n = 0.0, 0 + for bx, by in loader: + bx, by = bx.to(device), by.to(device) + out = model(bx) + bs = by.size(0) + total_mpe += calculate_mpjpe(out, by) * bs + pck = calculate_pck(out, by, thresholds=list(totals)) + for t in totals: + totals[t] += pck[t] * bs + n += bs + return {'samples': n, 'mpjpe': total_mpe / n, + **{f'pck@{int(t * 100)}': totals[t] / n for t in totals}} + + +def train_variant(v, dataset, device): + set_seed(SEED) + train_loader, val_loader, test_loader = create_preprocessed_train_val_test_loaders( + dataset=dataset, batch_size=BATCH, num_workers=2, random_seed=SEED) + + set_seed(SEED) # re-seed after split so init is split-independent + model = build_model(v).to(device) + info = describe(model) + print(f"[{v['name']}] params={info['params']:,} tcn_groups={info['tcn_groups_per_block']} " + f"conv_strides={info['conv_strides']} final_width={info['final_width']}", flush=True) + + criterion = PoseLoss(position_weight=1.0, bone_weight=0.2, loss_type='smooth_l1') + optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY, + betas=(0.9, 0.999)) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode='min', factor=0.5, patience=3, min_lr=LR / 1000, + cooldown=1, threshold=1e-4) + + best_val_mpe = float('inf') + best_val_pck20 = 0.0 + best_epoch = 0 + best_state = None + patience_counter = 0 + t0 = time.time() + error = None + epochs_run = 0 + + for epoch in range(1, EPOCHS + 1): + model.train() + ep_loss, nb = 0.0, 0 + te = time.time() + for i, (bx, by) in enumerate(train_loader): + bx = bx.to(device, non_blocking=True) + by = by.to(device, non_blocking=True) + optimizer.zero_grad(set_to_none=True) + out = model(bx) + loss, _parts = criterion(out, by) + if not torch.isfinite(loss): + error = f'non-finite loss at epoch {epoch} step {i}' + break + loss.backward() + optimizer.step() + ep_loss += loss.item() + nb += 1 + if epoch == 1 and i % 500 == 0: + print(f"[{v['name']}] e1 step {i}/{len(train_loader)} loss={loss.item():.5f}", + flush=True) + if error: + break + epochs_run = epoch + + val = evaluate(model, val_loader, device) + scheduler.step(val['mpjpe']) + lr_now = optimizer.param_groups[0]['lr'] + print(f"[{v['name']}] epoch {epoch}/{EPOCHS} train_loss={ep_loss / max(nb, 1):.5f} " + f"val_mpjpe={val['mpjpe']:.5f} val_pck20={val['pck@20'] * 100:.2f}% " + f"lr={lr_now:.2e} ({time.time() - te:.0f}s)", flush=True) + + if val['mpjpe'] < best_val_mpe: + best_val_mpe = val['mpjpe'] + best_val_pck20 = val['pck@20'] + best_epoch = epoch + best_state = copy.deepcopy(model.state_dict()) + patience_counter = 0 + else: + patience_counter += 1 + if patience_counter >= PATIENCE: + print(f"[{v['name']}] early stop at epoch {epoch} (best {best_epoch})", flush=True) + break + + train_seconds = time.time() - t0 + result = { + 'variant': v['name'], 'params': info['params'], + 'tcn_channels': v['tcn'], 'conv_channels': v['conv'], + 'attn_groups': v['attn_groups'], 'groups_mode': v['groups_mode'], + 'input_pw_groups': v['input_pw_groups'], + 'tcn_groups_per_block': info['tcn_groups_per_block'], + 'conv_strides': info['conv_strides'], 'final_width': info['final_width'], + 'batch_size': BATCH, 'max_epochs': EPOCHS, 'patience': PATIENCE, + 'lr': LR, 'weight_decay': WEIGHT_DECAY, 'seed': SEED, 'precision': 'fp32', + 'epochs_run': epochs_run, 'best_epoch': best_epoch, + 'best_val_mpjpe': best_val_mpe if best_state else None, + 'best_val_pck20': best_val_pck20 if best_state else None, + 'train_seconds': round(train_seconds, 1), + 'torch': torch.__version__, 'error': error, + 'finished_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), + } + + if best_state is not None: + ckpt = os.path.join(SWEEP, f"{v['name']}_best.pth") + torch.save(best_state, ckpt) + result['checkpoint'] = ckpt + model.load_state_dict(best_state) + + eval_loader = DataLoader(test_loader.dataset, batch_size=256, shuffle=False, + num_workers=2) + result['test_full'] = evaluate(model, eval_loader, device) + + w2f = dataset.window_to_file + clean_idx = [i for i in test_loader.dataset.indices if w2f[i] < CORRUPT_FILE_START] + clean_loader = DataLoader(Subset(dataset, clean_idx), batch_size=256, + shuffle=False, num_workers=2) + result['test_clean'] = evaluate(model, clean_loader, device) + print(f"[{v['name']}] TEST clean: pck20={result['test_clean']['pck@20'] * 100:.2f}% " + f"mpjpe={result['test_clean']['mpjpe']:.5f} | full: " + f"pck20={result['test_full']['pck@20'] * 100:.2f}%", flush=True) + return result + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('--dry-run', action='store_true', help='print param counts and exit') + args = ap.parse_args() + + if args.dry_run: + for v in VARIANTS: + m = build_model(v) + info = describe(m) + x = torch.randn(2, 540, 20) + m.eval() + y = m(x) + print(f"{v['name']:8s} params={info['params']:>9,} " + f"tcn={v['tcn']} conv={v['conv']} attn_g={v['attn_groups']} " + f"mode={v['groups_mode']} pw_g={v['input_pw_groups']} " + f"tcn_groups={info['tcn_groups_per_block']} strides={info['conv_strides']} " + f"W'={info['final_width']} out={tuple(y.shape)}") + return + + results_path = os.path.join(SWEEP, 'results.jsonl') + done = set() + if os.path.exists(results_path): + with open(results_path) as f: + for line in f: + try: + done.add(json.loads(line)['variant']) + except Exception: + pass + + device = torch.device('cuda') + print(f"torch {torch.__version__} on {torch.cuda.get_device_name(0)}", flush=True) + data_dir = os.path.join(BENCH, 'preprocessed_csi_data') + dataset = PreprocessedCSIKeypointsDataset(data_dir=data_dir, keypoint_scale=1000.0, + enable_temporal_clean=True) + + for v in VARIANTS: + if v['name'] in done: + print(f"[{v['name']}] already in results.jsonl — skipping", flush=True) + continue + print(f"\n===== variant: {v['name']} =====", flush=True) + try: + result = train_variant(v, dataset, device) + except Exception as e: # record and move on to next variant + import traceback + traceback.print_exc() + result = {'variant': v['name'], 'error': repr(e), + 'finished_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} + with open(results_path, 'a') as f: + f.write(json.dumps(result) + '\n') + f.flush() + print('\nSWEEP COMPLETE', flush=True) + + +if __name__ == '__main__': + main() diff --git a/benchmarks/wiflow-std/results/efficiency_sweep.jsonl b/benchmarks/wiflow-std/results/efficiency_sweep.jsonl new file mode 100644 index 00000000..99355884 --- /dev/null +++ b/benchmarks/wiflow-std/results/efficiency_sweep.jsonl @@ -0,0 +1,3 @@ +{"variant": "half", "params": 843834, "tcn_channels": [270, 220, 170, 120], "conv_channels": [4, 8, 16, 32], "attn_groups": 4, "groups_mode": "gcd20", "input_pw_groups": 1, "tcn_groups_per_block": [[20, 10], [10, 20], [20, 10], [10, 20]], "conv_strides": [2, 2, 2, 1], "final_width": 15, "batch_size": 64, "max_epochs": 50, "patience": 5, "lr": 0.0001, "weight_decay": 5e-05, "seed": 42, "precision": "fp32", "epochs_run": 28, "best_epoch": 23, "best_val_mpjpe": 0.008576328293592842, "best_val_pck20": 0.9690593021534107, "train_seconds": 1346.4, "torch": "2.11.0+cu128", "error": null, "finished_utc": "2026-06-11T03:09:47Z", "checkpoint": "/home/ruvultra/wiflow-std-bench/sweep/half_best.pth", "test_full": {"samples": 54000, "mpjpe": 0.009419974447676428, "pck@10": 0.8740543655289544, "pck@20": 0.9610469643628156, "pck@30": 0.9813556064146537, "pck@40": 0.9896086878246731, "pck@50": 0.9934827546013726}, "test_clean": {"samples": 52560, "mpjpe": 0.008980081718602137, "pck@10": 0.8840944136840205, "pck@20": 0.9662253179869514, "pck@30": 0.9847971080282144, "pck@40": 0.9917795997050618, "pck@50": 0.9946956242600532}} +{"variant": "quarter", "params": 338600, "tcn_channels": [135, 110, 85, 60], "conv_channels": [2, 4, 8, 16], "attn_groups": 2, "groups_mode": "gcd20", "input_pw_groups": 1, "tcn_groups_per_block": [[20, 5], [5, 10], [10, 5], [5, 20]], "conv_strides": [2, 2, 1, 1], "final_width": 15, "batch_size": 64, "max_epochs": 50, "patience": 5, "lr": 0.0001, "weight_decay": 5e-05, "seed": 42, "precision": "fp32", "epochs_run": 50, "best_epoch": 50, "best_val_mpjpe": 0.008780752391864856, "best_val_pck20": 0.9672531302240159, "train_seconds": 1754.4, "torch": "2.11.0+cu128", "error": null, "finished_utc": "2026-06-11T03:39:06Z", "checkpoint": "/home/ruvultra/wiflow-std-bench/sweep/quarter_best.pth", "test_full": {"samples": 54000, "mpjpe": 0.009705399298005634, "pck@10": 0.8646123917014511, "pck@20": 0.9553815319449813, "pck@30": 0.979827209190086, "pck@40": 0.9887037501511751, "pck@50": 0.9931309027671814}, "test_clean": {"samples": 52560, "mpjpe": 0.009279253277105465, "pck@10": 0.8742288637923323, "pck@20": 0.9605315079427745, "pck@30": 0.9833016723076865, "pck@40": 0.9908206971631566, "pck@50": 0.9942719799017071}} +{"variant": "tiny", "params": 56290, "tcn_channels": [68, 56, 44, 32], "conv_channels": [2, 4, 8, 16], "attn_groups": 2, "groups_mode": "depthwise", "input_pw_groups": 4, "tcn_groups_per_block": [[540, 68], [68, 56], [56, 44], [44, 32]], "conv_strides": [2, 1, 1, 1], "final_width": 16, "batch_size": 64, "max_epochs": 50, "patience": 5, "lr": 0.0001, "weight_decay": 5e-05, "seed": 42, "precision": "fp32", "epochs_run": 50, "best_epoch": 47, "best_val_mpjpe": 0.012602971208592256, "best_val_pck20": 0.9397210340146666, "train_seconds": 1540.1, "torch": "2.11.0+cu128", "error": null, "finished_utc": "2026-06-11T04:04:50Z", "checkpoint": "/home/ruvultra/wiflow-std-bench/sweep/tiny_best.pth", "test_full": {"samples": 54000, "mpjpe": 0.012859782406853305, "pck@10": 0.7640358444319831, "pck@20": 0.9364815320968628, "pck@30": 0.9731568422317505, "pck@40": 0.9866444962642811, "pck@50": 0.992488939108672}, "test_clean": {"samples": 52560, "mpjpe": 0.012502924276904246, "pck@10": 0.770895526488985, "pck@20": 0.9411073559313967, "pck@30": 0.9764840687790962, "pck@40": 0.9886695077067278, "pck@50": 0.9936238432039409}}