Wrapper code (speech-spoof-bench)

LICENSE.source

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Duc-Tuan Truong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

_net.py

@@ -0,0 +1,309 @@
+"""DPDA-GA XLS-R + TCM-Conformer network, ported from
+https://github.com/ductuantruong/dpda_ga (MIT) @ aab89614ee90252e10e5df35ec6aaa5ee2e179a7:
+  nets/conformer.py      (ConformerBlock + Head-Token Attention, verbatim below)
+  nets/conformer_tcm.py  (sinusoidal_embedding, MyConformer, SSLModel, Model)
+Hyper-parameters = main.py defaults (the README trains `--model conformer_tcm` with
+them): emb_size 144, heads 4, kernel_size 31, num_encoders 4 (ff_mult 4, exp_fac 2);
+verified against the released checkpoint shapes (strict load).
+
+Edits (inference-only, numerically identical; parameter/module names unchanged so
+the released state_dict loads strictly):
+  1. SSLModel builds the XLS-R 300M wav2vec2 architecture from the base
+     xlsr2_300m.pt model cfg merged onto the installed fairseq Wav2Vec2Config
+     (merge_with_parent; cross-version cfg-field fix, as W2V2-AASIST/_net.py),
+     instead of load_model_ensemble_and_task on a hard-coded path. The base
+     weights are not loaded -- the fine-tuned checkpoint overwrites every weight.
+  2. SSLModel.extract_feat: the `self.model.train()` call on a device/dtype
+     mismatch is removed (it would re-enable dropout at inference); the wrapper
+     moves the net to the device and calls .eval() itself.
+  3. Model.__init__ takes the hyper-parameters as keyword args (default = main.py)
+     instead of an argparse namespace; `device` arg dropped (unused at inference).
+  4. MyConformer.forward: the class-token prepend
+     `torch.stack([torch.vstack((cls, x[i])) for i in range(len(x))])` is written as
+     `torch.cat([cls.expand(B, 1, E), x], 1)` (same tensor; the python loop over the
+     batch would freeze the batch size into a traced ONNX graph).
+  5. Model.forward returns only the (B, 2) logits (attention maps dropped).
+  6. MyConformer.positional_emb is initialised to zeros instead of the slow python-list
+     sinusoid build: it is a frozen nn.Parameter stored in (and loaded from) the checkpoint.
+"""
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+from torch.nn.modules.transformer import _get_clones
+
+_XLSR_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "xlsr2_300m.pt")
+WINDOW = 66800  # source data_utils.Dataset_eval / OOD_Dataset_eval: self.cut = 66800
+
+
+def pad(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray:
+    """source utils.pad: first max_len samples; shorter clips repeat-tiled."""
+    x = np.asarray(x, dtype=np.float32).reshape(-1)
+    x_len = x.shape[0]
+    if x_len >= max_len:
+        return x[:max_len]
+    num_repeats = int(max_len / x_len) + 1
+    return np.tile(x, (1, num_repeats))[:, :max_len][0]
+
+
+# ============================ nets/conformer.py ============================
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def calc_same_padding(kernel_size):
+    pad = kernel_size // 2
+    return (pad, pad - (kernel_size + 1) % 2)
+
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * x.sigmoid()
+
+
+class GLU(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        out, gate = x.chunk(2, dim=self.dim)
+        return out * gate.sigmoid()
+
+
+class DepthWiseConv1d(nn.Module):
+    def __init__(self, chan_in, chan_out, kernel_size, padding):
+        super().__init__()
+        self.padding = padding
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in)
+
+    def forward(self, x):
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+class Scale(nn.Module):
+    def __init__(self, scale, fn):
+        super().__init__()
+        self.fn = fn
+        self.scale = scale
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+
+class Attention(nn.Module):
+    # Head Token attention: https://arxiv.org/pdf/2210.05958.pdf
+    def __init__(self, dim, heads=8, dim_head=64, qkv_bias=False, dropout=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = heads
+        inner_dim = dim_head * heads
+        self.scale = dim_head ** -0.5
+
+        self.qkv = nn.Linear(dim, inner_dim * 3, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(dropout)
+        self.proj = nn.Linear(inner_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.act = nn.GELU()
+        self.ht_proj = nn.Linear(dim_head, dim, bias=True)
+        self.ht_norm = nn.LayerNorm(dim_head)
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_heads, dim))
+
+    def forward(self, x, mask=None):
+        B, N, C = x.shape
+
+        # head token
+        head_pos = self.pos_embed.expand(x.shape[0], -1, -1)
+        x_ = x.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        x_ = x_.mean(dim=2)  # now the shape is [B, h, 1, d//h]
+        x_ = self.ht_proj(x_).reshape(B, -1, self.num_heads, C // self.num_heads)
+        x_ = self.act(self.ht_norm(x_)).flatten(2)
+        x_ = x_ + head_pos
+        x = torch.cat([x, x_], dim=1)
+
+        # normal mhsa
+        qkv = self.qkv(x).reshape(B, N + self.num_heads, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N + self.num_heads, C)
+        x = self.proj(x)
+
+        # merge head tokens into cls token
+        cls, patch, ht = torch.split(x, [1, N - 1, self.num_heads], dim=1)
+        cls = cls + torch.mean(ht, dim=1, keepdim=True) + torch.mean(patch, dim=1, keepdim=True)
+        x = torch.cat([cls, patch], dim=1)
+
+        x = self.proj_drop(x)
+
+        return x, attn
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult),
+            Swish(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class ConformerConvModule(nn.Module):
+    def __init__(self, dim, causal=False, expansion_factor=2, kernel_size=31, dropout=0.):
+        super().__init__()
+
+        inner_dim = dim * expansion_factor
+        padding = calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)
+
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Rearrange('b n c -> b c n'),
+            nn.Conv1d(dim, inner_dim * 2, 1),
+            GLU(dim=1),
+            DepthWiseConv1d(inner_dim, inner_dim, kernel_size=kernel_size, padding=padding),
+            nn.BatchNorm1d(inner_dim) if not causal else nn.Identity(),
+            Swish(),
+            nn.Conv1d(inner_dim, dim, 1),
+            Rearrange('b c n -> b n c'),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class ConformerBlock(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8, ff_mult=4, conv_expansion_factor=2,
+                 conv_kernel_size=31, attn_dropout=0., ff_dropout=0., conv_dropout=0.,
+                 conv_causal=False):
+        super().__init__()
+        self.ff1 = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+        self.attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout)
+        self.conv = ConformerConvModule(dim=dim, causal=conv_causal, expansion_factor=conv_expansion_factor,
+                                        kernel_size=conv_kernel_size, dropout=conv_dropout)
+        self.ff2 = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+
+        self.attn = PreNorm(dim, self.attn)
+        self.ff1 = Scale(0.5, PreNorm(dim, self.ff1))
+        self.ff2 = Scale(0.5, PreNorm(dim, self.ff2))
+
+        self.post_norm = nn.LayerNorm(dim)
+
+    def forward(self, x, mask=None):
+        x = self.ff1(x) + x
+        attn_x, attn_weight = self.attn(x, mask=mask)
+        x = attn_x + x
+        x = self.conv(x) + x
+        x = self.ff2(x) + x
+        x = self.post_norm(x)
+        return x, attn_weight
+
+
+# ============================ nets/conformer_tcm.py ============================
+def sinusoidal_embedding(n_channels, dim):
+    pe = torch.FloatTensor([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
+                            for p in range(n_channels)])
+    pe[:, 0::2] = torch.sin(pe[:, 0::2])
+    pe[:, 1::2] = torch.cos(pe[:, 1::2])
+    return pe.unsqueeze(0)
+
+
+class MyConformer(nn.Module):
+    def __init__(self, emb_size=128, heads=4, ffmult=4, exp_fac=2, kernel_size=16, n_encoders=1):
+        super(MyConformer, self).__init__()
+        self.dim_head = int(emb_size / heads)
+        self.dim = emb_size
+        self.heads = heads
+        self.kernel_size = kernel_size
+        self.n_encoders = n_encoders
+        # edit 6: zeros placeholder; the checkpoint stores this (frozen) table and overwrites it
+        self.positional_emb = nn.Parameter(torch.zeros(1, 10000, emb_size), requires_grad=False)
+        self.encoder_blocks = _get_clones(ConformerBlock(dim=emb_size, dim_head=self.dim_head, heads=heads,
+                                                         ff_mult=ffmult, conv_expansion_factor=exp_fac,
+                                                         conv_kernel_size=kernel_size),
+                                          n_encoders)
+        self.class_token = nn.Parameter(torch.rand(1, emb_size))
+        self.fc5 = nn.Linear(emb_size, 2)
+
+    def forward(self, x):  # x shape [bs, tiempo, frecuencia]
+        x = x + self.positional_emb[:, :x.size(1), :]
+        # edit 4: == torch.stack([torch.vstack((self.class_token, x[i])) for i in range(len(x))])
+        x = torch.cat([self.class_token.unsqueeze(0).expand(x.size(0), -1, -1), x], dim=1)  # [bs,1+tiempo,emb]
+        for layer in self.encoder_blocks:
+            x, _ = layer(x)  # [bs,1+tiempo,emb_size]
+        embedding = x[:, 0, :]  # [bs, emb_size]
+        out = self.fc5(embedding)  # [bs,2]
+        return out
+
+
+class SSLModel(nn.Module):  # W2V
+    def __init__(self, cp_path=_XLSR_PATH):
+        super(SSLModel, self).__init__()
+        import fairseq
+        from fairseq.dataclass.utils import merge_with_parent
+        from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model
+        state = fairseq.checkpoint_utils.load_checkpoint_to_cpu(cp_path)
+        cfg = merge_with_parent(Wav2Vec2Config(), state["cfg"]["model"], remove_missing=True)
+        del state
+        self.model = Wav2Vec2Model.build_model(cfg, task=None)
+        self.out_dim = 1024
+
+    def extract_feat(self, input_data):
+        # edit 2: no `self.model.to(...); self.model.train()` here.
+        if input_data.ndim == 3:
+            input_tmp = input_data[:, :, 0]
+        else:
+            input_tmp = input_data
+        # [batch, length, dim]
+        emb = self.model(input_tmp, mask=False, features_only=True)['x']
+        return emb
+
+
+class Model(nn.Module):
+    def __init__(self, emb_size=144, heads=4, kernel_size=31, num_encoders=4):
+        super().__init__()
+        self.ssl_model = SSLModel()
+        self.LL = nn.Linear(1024, emb_size)
+        self.first_bn = nn.BatchNorm2d(num_features=1)
+        self.selu = nn.SELU(inplace=True)
+        self.conformer = MyConformer(emb_size=emb_size, n_encoders=num_encoders,
+                                     heads=heads, kernel_size=kernel_size)
+
+    def forward(self, x):
+        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
+        x = self.LL(x_ssl_feat)  # (bs,frame_number,feat_out_dim)
+        x = x.unsqueeze(dim=1)  # add channel
+        x = self.first_bn(x)
+        x = self.selu(x)
+        x = x.squeeze(dim=1)
+        out = self.conformer(x)
+        return out

dpda_ga_tcm.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+from speech_spoof_bench.model import AntiSpoofingModel
+
+from _net import WINDOW, Model as DPDAGATCMNet, pad
+
+_DIR = os.path.dirname(os.path.abspath(__file__))
+# Elementwise mean of HF tuantd/dpda_ga pcgrad/conformer_tcm/best/best_{0..3}.pth (make_avg.py);
+# best_4.pth on HF is truncated, so the authors' 5-checkpoint average is not reproducible.
+_CKPT = os.path.join(_DIR, "dpda_ga_tcm_avg4.pth")
+
+
+class DPDAGATCM(AntiSpoofingModel):
+    name = "DPDA-GA-TCM"
+    expected_sample_rate = 16000
+    batch_size = 24  # sweep 2026-10-30 (cudnn TF32 off): RTX 4070 Ti SUPER 121.7 utt/s @24 (plateau 16-64); RTX 5060 Ti -> 4 (see load)
+
+    def load(self) -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # Strict FP32: cuDNN's default TF32 convs (XLS-R CNN front-end) make scores batch-dependent
+        # (~2e-4); off -> batched == single within 2e-6. No speed cost on the 4070 Ti SUPER.
+        torch.backends.cudnn.allow_tf32 = False
+        net = DPDAGATCMNet()  # emb 144, heads 4, kernel 31, 4 encoders (main.py defaults)
+        sd = torch.load(_CKPT, map_location="cpu")
+        net.load_state_dict(sd, strict=True)
+        # .to() before .eval(); extract_feat no longer flips XLS-R back to train().
+        self.net = net.to(self.device).eval()
+        if self.device == "cuda" and "5060" in torch.cuda.get_device_name(0):
+            self.batch_size = 4  # sweep 2026-10-30: RTX 5060 Ti peaks at bs 4 (22.2 utt/s), ~15 at bs >= 8
+
+    @torch.no_grad()
+    def score_batch(self, audios, srs):
+        x = np.stack([pad(a, WINDOW) for a in audios])
+        logits = self.net(torch.from_numpy(x).to(self.device)).float().reshape(-1, 2)
+        # source main.py produce_evaluation_file: score = batch_out[:, 1]; label 1 == bonafide
+        return logits[:, 1].cpu().tolist()
+
+    def unload(self) -> None:
+        self.net = None
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

implementation-notes.md

@@ -0,0 +1,58 @@
+# DPDA-GA-TCM — implementation notes
+
+Plan: `speech-spoof-bench/docs/plans/2026-10-29-mamba-batch-arena-submission.md` (model #2).
+
+## Source
+- Code: github.com/ductuantruong/dpda_ga @ `aab89614ee90252e10e5df35ec6aaa5ee2e179a7` (MIT, `LICENSE.source`):
+  `nets/conformer.py` + `nets/conformer_tcm.py` → `_net.py` (edits listed in its header).
+- Weights: HF `tuantd/dpda_ga` @ `531554e0a1c28d6a6ccf22a355281e3d3acefeb5`,
+  `pcgrad/conformer_tcm/best/best_{0..4}.pth` (downloaded by `benchmarks/_ckpt_dl/dl.py`).
+- Hyper-parameters = `main.py` defaults: emb 144, heads 4, kernel 31, 4 encoders; checkpoint
+  shapes confirm (LL 144×1024, ht_proj 144×36, depthwise conv k=31, encoder_blocks 0..3).
+- XLS-R 300M architecture built from `xlsr2_300m.pt` cfg (symlink to `../W2V2-AASIST/`) via
+  `merge_with_parent`; base weights not loaded (the checkpoint overwrites all 588 tensors).
+
+## Checkpoint averaging (`make_avg.py` → `dpda_ga_tcm_avg4.pth`)
+- `best_4.pth` is truncated on HF: 547,680,256 bytes vs 1,285,107,746 for the others, no ZIP
+  central directory; `torch.load` → "failed finding central directory". Confirmed locally.
+- `main.py` eval averages `n_average_model` (default 5) n-best checkpoints, but
+  `sd = model.state_dict()` aliases the live parameters, so the following
+  `model.load_state_dict(best_1)` overwrites `sd` in place. Net effect:
+  `(2*b1 + b2 + … + b_{n-1}) / n`; best_0 never contributes. Maintainer decision (2026-10-30):
+  reproduce main.py exactly with the 4 intact files → `(2*best_1 + best_2 + best_3) / 4`,
+  float32, main.py's summation order; verified bit-identical to literally running main.py's
+  loop on a module holding the tensors. Int buffers (num_batches_tracked) get the truncated
+  float quotient, as main.py's copy into a long buffer does (unused at eval).
+- The paper's numbers come from the 5-file version `(2*b1+b2+b3+b4)/5`, so small deviations
+  from the paper are expected.
+
+## Wrapper decisions
+- Input: first 66,800 samples; shorter clips repeat-tiled (`utils.pad`). No resampling/normalisation.
+- Score = `logits[:, 1]` (label 1 = bonafide in `utils.read_metadata`; `produce_evaluation_file`).
+- `.to(device)` then `.eval()`; `extract_feat`'s `self.model.train()` removed.
+- cuDNN TF32 disabled in `load()`: with it on, batched vs single scores differed by up to 2e-4
+  (XLS-R conv front-end); off → 1.7e-6. 4070 Ti SUPER speed unchanged (121.7 utt/s @24 both ways);
+  5060 Ti PyTorch drops 47.5 → 22.2 utt/s (bs 4), but that card runs TensorRT.
+- Params: 321,213,220 (incl. the frozen 10000×144 positional table, 1.44 M).
+
+## Batch sweep (sweep.py, 66,800-sample random input, 2026-10-30)
+| card | TF32 off best | TF32 on best |
+|---|---|---|
+| RTX 4070 Ti SUPER (idx 2) | 121.7 utt/s @24 | 121.7 @24 |
+| RTX 5060 Ti (idx 3) | 22.2 @4 | 47.5 @4 |
+Raw: `logs_sweep_g{2,3}_tf{0,1}.txt`.
+
+## Sanity (full J-SPAW_LA, PyTorch, RTX 4070 Ti SUPER)
+- EER 12.90% (2397 trials, 0 skipped), ~110 utt/s. Comparable XLS-R models: DPDA-GA-AASIST 13.65,
+  XLSR-SLS 8.70, W2V2-AASIST 26.17. Not flipped.
+
+## TensorRT (`trt_dpda_ga_tcm.py`, sidecar `trt_dpda-ga-tcm.json`, TRT 10.16)
+Adapted from XLSR-MamBo's script (config block only). The class-token loop rewrite (`_net.py`
+edit 4) keeps the batch axis dynamic in the traced ONNX.
+| card | precision | best batch | utt/s (sweep) | build parity Spearman (64 ITW) | J-SPAW EER | J-SPAW Spearman vs PT / max Δ |
+|---|---|---|---|---|---|---|
+| RTX 4070 Ti SUPER (idx 2) | FP32 (FP16 Spearman 0.99963 < 0.9999) | 8 | 173.3 | 1.0000 | 12.8992 (=PT) | 0.999998 / 0.031 |
+| RTX 5060 Ti (idx 3) | FP16 | 16 | 362.2 | 0.99991 | 12.8992 (=PT) | 0.999917 / 0.287 |
+Both cards pass (|ΔEER| = 0 ≤ 0.25, Spearman ≥ 0.999). J-SPAW_LA is submitted from PyTorch scores
+(`results_pt/`); TRT runs in `results_trt/` (idx 2) and `results_trt3/` (idx 3). `cmp_scores.py`
+does the comparison. The unused 4070 FP16 engine was deleted.

make_avg.py

@@ -0,0 +1,47 @@
+# make_avg.py — build the averaged DPDA-GA-TCM checkpoint (2026-10-29).
+# Source: HF tuantd/dpda_ga @ 531554e0a1c28d6a6ccf22a355281e3d3acefeb5, pcgrad/conformer_tcm/best/best_{0..4}.pth.
+# best_4.pth on HF is truncated (547,680,256 bytes vs ~1,285,107,746; no ZIP central directory,
+# torch.load fails), so we use the 4 intact n-best checkpoints best_0..best_3 (n_average_model 5 -> 4)
+# and reproduce main.py's averaging EXACTLY, including its aliasing bug (maintainer decision
+# 2026-10-30): `sd = model.state_dict()` returns tensors that share storage with the model, so the
+# next `model.load_state_dict(best_1)` overwrites sd in place; the loop then adds best_1 to itself.
+# main.py's result is therefore (2*b1 + b2 + ... + b_{n-1}) / n -- best_0 never contributes.
+# For n = 4: (2*best_1 + best_2 + best_3) / 4, computed in float32 in main.py's order. Integer
+# buffers (BatchNorm num_batches_tracked, unused at eval) get main.py's float quotient truncated
+# by the copy into the long buffer.
+import os
+import sys
+
+import torch
+
+SRC = sys.argv[1] if len(sys.argv) > 1 else os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "..", "_ckpt_dl", "tuantd__dpda_ga", "pcgrad", "conformer_tcm", "best")
+OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dpda_ga_tcm_avg4.pth")
+N = 4
+
+try:
+    torch.load(os.path.join(SRC, "best_4.pth"), map_location="cpu")
+    raise SystemExit("best_4.pth loads -- average all 5 instead")
+except RuntimeError as e:
+    print("best_4.pth unreadable as expected:", str(e)[:100])
+
+sds = [torch.load(os.path.join(SRC, f"best_{i}.pth"), map_location="cpu") for i in range(N)]
+keys = list(sds[0].keys())
+for i, sd in enumerate(sds[1:], 1):
+    assert list(sd.keys()) == keys, f"key mismatch in best_{i}"
+    for k in keys:
+        assert sd[k].shape == sds[0][k].shape and sd[k].dtype == sds[0][k].dtype, (i, k)
+avg = {}
+for k in keys:
+    t0 = sds[0][k]
+    acc = sds[1][k].clone()          # sd aliased to the params after load_state_dict(best_1)
+    acc = acc + sds[1][k]            # i = 1: sd[key] = sd[key] + sd2[key]  -> 2*b1
+    for i in range(2, N):
+        acc = acc + sds[i][k]        # i = 2..N-1
+    acc = acc / N                    # true division (float even for int buffers)
+    avg[k] = acc.to(t0.dtype)        # load_state_dict copy_ into the original dtype (long: trunc)
+dtypes = sorted({str(v.dtype) for v in avg.values()})
+print(f"{len(keys)} keys, dtypes {dtypes}; identical ckpts? ",
+      all(torch.equal(sds[0][k], sds[1][k]) for k in keys))
+torch.save(avg, OUT)
+print("wrote", OUT, os.path.getsize(OUT))

meta.yaml

@@ -0,0 +1,45 @@
+system:
+  name: "DPDA-GA-TCM"
+  slug: "dpda-ga-tcm"
+  description: >
+    DPDA-GA (dual-path data-augmented training with gradient alignment) applied to an
+    XLS-R 300M + TCM-Conformer detector: a fully fine-tuned wav2vec 2.0 XLS-R 300M
+    front-end projected to 144-d (BatchNorm + SELU), followed by 4 Conformer blocks with
+    Head-Token (TCM) multi-head attention (4 heads, depthwise conv kernel 31) and a class
+    token feeding a linear 2-way head. Trained on ASVspoof2019 LA with each batch seen both
+    clean and RawBoost-augmented, the two losses' gradients aligned with PCGrad.
+    Checkpoint = weight average of the authors' n-best checkpoints, reproducing their
+    main.py averaging exactly: the release's 5th checkpoint (best_4.pth) is truncated on
+    HuggingFace and cannot be loaded, so 4 checkpoints are averaged instead of 5, and
+    main.py's averaging code drops best_0 and double-counts best_1, giving
+    (2*best_1 + best_2 + best_3)/4. FP32, deterministic first-66800-sample window
+    (shorter clips repeat-tiled).
+  code: "https://github.com/ductuantruong/dpda_ga"
+  checkpoint: "https://huggingface.co/SpeechAntiSpoofingBenchmarks/DPDA-GA-TCM"
+  params_millions: 321.2132
+  paper:
+    arxiv_id: "2509.20682"
+    url: "https://arxiv.org/abs/2509.20682"
+    bibtex: |
+      @inproceedings{truong2026dpdaga,
+        author    = {Truong, Duc-Tuan and Liu, Tianchi and Li, Junjie and Tao, Ruijie and Lee, Kong Aik and Chng, Eng Siong},
+        title     = {Addressing Gradient Misalignment in Data-Augmented Training for Robust Speech Deepfake Detection},
+        booktitle = {ICASSP 2026 - 2026 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+        year      = {2026},
+        pages     = {16252--16256},
+        doi       = {10.1109/ICASSP55912.2026.11464739}
+      }
+notes: >
+  DPDA-GA XLS-R + TCM-Conformer (PCGrad) from ductuantruong/dpda_ga @ aab8961 (MIT), weights
+  from HF tuantd/dpda_ga @ 531554e pcgrad/conformer_tcm/best/best_{0..4}.pth. best_4.pth is
+  truncated on HF (547,680,256 of ~1,285,107,746 bytes; no ZIP central directory), so the
+  authors' 5-checkpoint average (main.py --average_model True --n_average_model 5) cannot be
+  rebuilt; we apply main.py's averaging to the 4 intact checkpoints. main.py's loop has a
+  state_dict aliasing bug (sd = model.state_dict() shares storage with the model, so loading
+  best_1 overwrites it): its result is (2*b1 + b2 + ... + b_{n-1})/n with best_0 dropped. We
+  reproduce it exactly: (2*best_1 + best_2 + best_3)/4 (make_avg.py). Strict load.
+  Deterministic first-66800-sample window, shorter clips repeat-tiled (source utils.pad), no
+  resampling/normalisation. score = logits[:, 1] (bona fide = class 1; source
+  produce_evaluation_file); higher = more bona fide. FP32 with cuDNN TF32 off (TF32 convs made
+  scores batch-dependent by ~2e-4). Paper EERs (5-checkpoint model): 21DF 1.81%, ITW 6.48%.
+  Trained on ASVspoof2019 LA train (dev for model selection) only.

test_dpda_ga_tcm.py

@@ -0,0 +1,54 @@
+import numpy as np
+import pytest
+import torch
+
+import _net
+import dpda_ga_tcm as D
+
+N_PARAMS = 321_213_220  # incl. the frozen 10000x144 positional table (1.44 M)
+
+
+def test_pad_long_truncated():
+    x = np.arange(100000, dtype=np.float32)
+    assert np.array_equal(_net.pad(x), x[:66800])
+
+
+def test_pad_short_repeated():
+    x = np.arange(1000, dtype=np.float32)
+    out = _net.pad(x)
+    assert out.shape == (66800,) and np.array_equal(out[:1000], x) and np.array_equal(out[1000:2000], x)
+
+
+@pytest.fixture(scope="module")
+def model():
+    m = D.DPDAGATCM(); m.load()   # strict=True load
+    yield m
+    m.unload()
+
+
+def test_eval_mode_and_params(model):
+    assert not any(mod.training for mod in model.net.modules())
+    assert sum(p.numel() for p in model.net.parameters()) == N_PARAMS
+
+
+def test_cls_token_cat_matches_source_stack(model):
+    c = model.net.conformer
+    x = torch.randn(3, 10, 144, device=model.device)
+    src = torch.stack([torch.vstack((c.class_token, x[i])) for i in range(len(x))])
+    new = torch.cat([c.class_token.unsqueeze(0).expand(3, -1, -1), x], dim=1)
+    assert torch.equal(src, new)
+
+
+def test_batch_matches_single(model):
+    rng = np.random.default_rng(0)
+    a = [rng.standard_normal(n).astype(np.float32) * s for n, s in ((48000, 0.1), (20000, 0.5), (70000, 0.02), (3000, 0.3))]
+    batched = model.score_batch(a, [16000] * 4)
+    single = [model.score_batch([x], [16000])[0] for x in a]
+    assert len(batched) == 4 and all(np.isfinite(batched))
+    assert np.allclose(batched, single, atol=1e-4)
+
+
+def test_deterministic_and_finite(model):
+    a = [np.random.randn(40000).astype(np.float32) for _ in range(3)]
+    s1 = model.score_batch(a, [16000] * 3)
+    assert all(np.isfinite(s1)) and s1 == model.score_batch(a, [16000] * 3)

trt_dpda_ga_tcm.py

@@ -0,0 +1,569 @@
+#!/usr/bin/env python3
+"""TensorRT export + inference for DPDAGATCM (DPDA-GA XLS-R + TCM conformer). Self-contained (no shared package).
+
+Exports only the model's `net` (all preprocessing already lives in the original
+`score_batch`) with a fixed time axis and a dynamic batch axis, builds a FP16
+engine (FP32 fallback if parity drifts), finds the fastest batch on the current
+GPU, and exposes a drop-in `DPDAGATCMTRT` class identical to the PyTorch path except
+the neural forward runs on TensorRT.
+
+CLI:
+  python trt_dpda_ga_tcm.py export   # ONNX -> engine -> parity -> sweep -> sidecar
+  python trt_dpda_ga_tcm.py sweep     # re-run the batch sweep, update sidecar
+  python trt_dpda_ga_tcm.py parity    # PyTorch vs TRT parity report
+  python trt_dpda_ga_tcm.py score AUDIO.wav
+
+Pin the GPU with:  CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=<n>
+"""
+from __future__ import annotations
+import argparse
+import io
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))                       # import dir-local entry + _net
+# Pin GPU deterministically: PCI order makes CUDA indices match `nvidia-smi`.
+os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", os.environ.get("SSB_TRT_GPU", "2"))
+
+import torch  # noqa: E402  (after env pin)
+
+# ======================= per-model config =======================
+ENTRY_MODULE   = "dpda_ga_tcm"          # module exposing the AntiSpoofingModel subclass
+ENTRY_CLASS    = "DPDAGATCM"     # the subclass name
+SLUG           = "dpda-ga-tcm"
+PARITY_DATASET = "InTheWild"       # sibling dataset dir with data/*.parquet
+MAX_BATCH_CAP  = 32                # VRAM ceiling for the profile + sweep
+PARITY_CHUNK   = 8                 # safe mini-batch for the parity comparison
+OPSET          = 17
+# Keep FP16 iff it preserves the score RANKING (Spearman) -> identical EER.
+# This is the metric that matters for the benchmark and is scale-invariant, so
+# small absolute-logit drift (harmless for EER) does not force an FP32 fallback.
+# FP16 is also mandatory for the largest models (FP32 would not fit in VRAM).
+PARITY_SPEARMAN_TOL = 0.9999       # min Spearman rank-corr to keep FP16
+PARITY_FLOOR        = 0.99         # hard floor: below this the engine is wrong -> FAIL
+PARITY_MAD_TOL = 1e-2              # informational only
+PARITY_R_TOL   = 0.9999           # informational only
+FORCE_FP32         = bool(int(os.environ.get("TRT_FORCE_FP32", "0")))
+NO_TF32            = bool(int(os.environ.get("TRT_NO_TF32", "0")))  # per-GPU choice, see build_engine
+FORCE_FP16         = False   # skip FP32 (for giant models where FP32 won't fit VRAM)
+DYNAMO_EXPORT      = False   # 1.28 GB < 2 GB protobuf limit
+ALLOW_ORT_FALLBACK = False
+# ================================================================
+
+from importlib import import_module as _imp  # noqa: E402
+_OrigClass = getattr(_imp(ENTRY_MODULE), ENTRY_CLASS)
+
+
+# ----------------------------------------------------------------------------
+# helpers
+# ----------------------------------------------------------------------------
+def gpu_slug() -> str:
+    name = torch.cuda.get_device_name(0)
+    return name.replace("NVIDIA ", "").replace("GeForce ", "").strip().replace(" ", "_")
+
+
+def load_model():
+    m = _OrigClass()
+    m.load()
+    return m
+
+
+def real_audio(n=64):
+    """Decode up to n real 16 kHz mono utterances from PARITY_DATASET/data/*.parquet."""
+    import pyarrow.parquet as pq
+    import soundfile as sf
+    import torchaudio.functional as AF
+
+    data_dir = HERE.parent / PARITY_DATASET / "data"
+    files = sorted(data_dir.glob("test-*.parquet")) or sorted(data_dir.glob("*.parquet"))
+    out = []
+    for f in files:
+        t = pq.read_table(f)
+        col = "audio" if "audio" in t.column_names else t.column_names[0]
+        for row in t.column(col).to_pylist():
+            b = row["bytes"] if isinstance(row, dict) else row
+            if not b:
+                continue
+            a, sr = sf.read(io.BytesIO(b), dtype="float32")
+            if a.ndim > 1:
+                a = a.mean(1)
+            a = np.ascontiguousarray(a, dtype=np.float32)
+            if sr != 16000:
+                a = AF.resample(torch.from_numpy(a), sr, 16000).numpy().astype(np.float32)
+                sr = 16000
+            out.append((a, sr))
+            if len(out) >= n:
+                return out
+    if not out:
+        raise RuntimeError(f"no parity audio found under {data_dir}")
+    return out
+
+
+class _Capture:
+    """Wrap net: pass through to the real net, record input tensor + output."""
+
+    def __init__(self, net):
+        self.net = net
+        self.x = None
+        self.out = None
+
+    def __call__(self, x, *a, **k):
+        self.x = x.detach()
+        self.out = self.net(x, *a, **k)
+        return self.out
+
+    def __getattr__(self, name):
+        return getattr(self.net, name)
+
+
+def _logits_index(out):
+    """Return (L, i, n_classes): tuple length (None if tensor), logits slot, n_classes.
+
+    Heuristic: the class-logits tensor is the 2-D (B, C) tensor with the smallest C.
+    """
+    if isinstance(out, torch.Tensor):
+        return None, None, int(out.shape[-1])
+    cands = [(j, t) for j, t in enumerate(out)
+             if isinstance(t, torch.Tensor) and t.dim() == 2]
+    if not cands:
+        raise RuntimeError("could not locate a 2-D (B,C) logits tensor in net output")
+    j, t = min(cands, key=lambda it: int(it[1].shape[-1]))
+    return len(out), j, int(t.shape[-1])
+
+
+def analyze(model):
+    """One real forward through the capture shim -> (T, L, i, n_classes)."""
+    data = real_audio(1)
+    audios = [a for a, _ in data]
+    srs = [s for _, s in data]
+    cap = _Capture(model.net)
+    model.net = cap
+    with torch.no_grad():
+        model.score_batch(audios, srs)
+    model.net = cap.net
+    T = int(cap.x.shape[-1])
+    L, i, n_classes = _logits_index(cap.out)
+    return T, L, i, n_classes
+
+
+def _extractor(L, i):
+    """Pick the logits tensor out of a net's raw output."""
+    if L is None:
+        return lambda y: y
+    return lambda y, i=i: y[i]
+
+
+def _rebuild(L, i):
+    """Wrap a bare logits tensor back into the net's original output structure."""
+    if L is None:
+        return lambda y: y
+    return lambda y, L=L, i=i: tuple(y if j == i else None for j in range(L))
+
+
+def _prep_for_export(net):
+    """Make export-hostile layers traceable. No-op for non-fairseq models.
+
+    fairseq wav2vec2/hubert call `pad_to_multiple`, which does `(tsz/multiple)
+    .is_integer()`; under torch.jit tracing `tsz` becomes a Tensor with no
+    `.is_integer()`. Our time axis is static, so we swap in a constant-length
+    pad that traces cleanly. Patches every fairseq module that bound the name.
+    """
+    def _safe_pad(x, multiple, dim=-1, value=0):
+        import torch.nn.functional as F
+        if x is None:
+            return None, 0
+        tsz = int(x.shape[dim])                      # static: time axis is fixed
+        rem = (multiple - tsz % multiple) % multiple
+        if rem == 0:
+            return x, 0
+        pad_offset = (0,) * (-1 - dim) * 2
+        return F.pad(x, (*pad_offset, 0, rem), value=value), rem
+
+    for modname in ("fairseq.models.wav2vec.utils",
+                    "fairseq.models.wav2vec.wav2vec2",
+                    "fairseq.models.hubert.hubert"):
+        mod = sys.modules.get(modname)
+        if mod is not None and hasattr(mod, "pad_to_multiple"):
+            mod.pad_to_multiple = _safe_pad
+    _freeze_sinc(net)
+    # optional per-model export patch (dir-local `_trt_patch.py` with `patch(net)`)
+    try:
+        import importlib
+        importlib.import_module("_trt_patch").patch(net)
+    except ModuleNotFoundError:
+        pass
+    if DYNAMO_EXPORT:
+        _replace_global_avgpool(net)
+    return net
+
+
+class _MeanPool(torch.nn.Module):
+    """Global average over `dims` (keepdim) — == AdaptiveAvgPool{1,2}d(1)."""
+
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+
+    def forward(self, x):
+        return x.mean(dim=self.dims, keepdim=True)
+
+
+def _replace_global_avgpool(net):
+    """Swap AdaptiveAvgPool1d/2d(output_size=1) for an explicit mean. The dynamo
+    exporter lowers the adaptive pool to as_strided/SequenceEmpty, which TensorRT
+    rejects; a plain mean lowers to ReduceMean. Identical for output_size==1."""
+    import torch.nn as nn
+    for full_name, mod in list(net.named_modules()):
+        is1d = isinstance(mod, nn.AdaptiveAvgPool1d) and mod.output_size in (1, (1,))
+        is2d = isinstance(mod, nn.AdaptiveAvgPool2d) and mod.output_size in (1, (1, 1))
+        if not (is1d or is2d):
+            continue
+        parent = net
+        *parents, attr = full_name.split(".")
+        for p in parents:
+            parent = getattr(parent, p)
+        setattr(parent, attr, _MeanPool((-1,) if is1d else (-2, -1)))
+    return net
+
+
+def _freeze_sinc(net):
+    """Replace SincConv-style layers with an equivalent nn.Conv1d holding the
+    precomputed band-pass filters. At eval the filters are constant, but their
+    in-forward construction (torch.sin/cat/flip from learnable params) either
+    won't build in TensorRT or constant-folds to wrong values. Baking them into a
+    plain Conv1d removes the sinc math from the graph. No-op when no Sinc layer.
+    """
+    import torch.nn as nn
+    sincs = [(n, m) for n, m in net.named_modules() if "Sinc" in type(m).__name__]
+    if not sincs:
+        return net
+    dev = next(net.parameters()).device
+    for full_name, mod in sincs:
+        kernel = int(getattr(mod, "kernel_size", 0)) or 1
+        with torch.no_grad():
+            try:
+                mod(torch.zeros(1, 1, max(kernel * 4, 4096), device=dev))
+            except Exception:  # noqa: BLE001 — filters are set before the conv call
+                pass
+        W = mod.filters.detach().clone()          # [out, 1, kernel] (or [out, kernel])
+        if W.dim() == 2:
+            W = W.unsqueeze(1)
+        conv = nn.Conv1d(W.shape[1], W.shape[0], W.shape[2],
+                         stride=int(getattr(mod, "stride", 1)),
+                         padding=int(getattr(mod, "padding", 0)),
+                         dilation=int(getattr(mod, "dilation", 1)),
+                         bias=False).to(dev).eval()
+        conv.weight.data.copy_(W)
+        parent = net
+        *parents, attr = full_name.split(".")
+        for p in parents:
+            parent = getattr(parent, p)
+        setattr(parent, attr, conv)
+    return net
+
+
+class _ExportNet(torch.nn.Module):
+    """forward(x[B,T]) -> logits[B,C] (single tensor) for ONNX/TRT."""
+
+    def __init__(self, net, L, i):
+        super().__init__()
+        self.net = net
+        self._extract = _extractor(L, i)
+
+    def forward(self, x):
+        return self._extract(self.net(x))  # Model takes (B,T)
+
+
+# ----------------------------------------------------------------------------
+# export + build
+# ----------------------------------------------------------------------------
+def export_onnx(model, T, L, i, onnx_path):
+    net = _prep_for_export(model.net)
+    wrap = _ExportNet(net, L, i).eval().to("cuda")
+    dummy = torch.zeros(2, T, device="cuda", dtype=torch.float32)
+    if DYNAMO_EXPORT:
+        # >2 GB models: TorchScript exporter's shape-inference overflows the 2 GB
+        # protobuf limit. The dynamo exporter writes weights as external data.
+        batch = torch.export.Dim("b", min=1, max=MAX_BATCH_CAP)
+        torch.onnx.export(
+            wrap, (dummy,), str(onnx_path), dynamo=True, external_data=True,
+            input_names=["wav"], output_names=["logits"],
+            dynamic_shapes={"x": {0: batch}},
+        )
+    else:
+        torch.onnx.export(
+            wrap, dummy, str(onnx_path), opset_version=OPSET,
+            input_names=["wav"], output_names=["logits"],
+            dynamic_axes={"wav": {0: "batch"}, "logits": {0: "batch"}},
+            do_constant_folding=True,
+        )
+    return onnx_path
+
+
+def build_engine(onnx_path, T, precision, max_batch, opt_batch, engine_path, timing_cache):
+    import tensorrt as trt
+
+    sev = trt.Logger.VERBOSE if os.environ.get("SSB_TRT_VERBOSE") else trt.Logger.WARNING
+    logger = trt.Logger(sev)
+    builder = trt.Builder(logger)
+    network = builder.create_network(0)
+    parser = trt.OnnxParser(network, logger)
+    # parse_from_file resolves external-data sidecars (needed for >2 GB models);
+    # works for inline ONNX too.
+    if not parser.parse_from_file(str(onnx_path)):
+        errs = "; ".join(str(parser.get_error(k)) for k in range(parser.num_errors))
+        raise RuntimeError(f"onnx parse failed: {errs}")
+
+    cfg = builder.create_builder_config()
+    cfg.builder_optimization_level = 1                       # minimum build time
+    if precision == "fp16":
+        cfg.set_flag(trt.BuilderFlag.FP16)
+    if NO_TF32:
+        # strict FP32: the RTX 5060 Ti TF32 engine drifted (J-SPAW Spearman 0.9969 vs PyTorch)
+        cfg.clear_flag(trt.BuilderFlag.TF32)
+
+    tc_bytes = Path(timing_cache).read_bytes() if Path(timing_cache).exists() else b""
+    tc = cfg.create_timing_cache(tc_bytes)
+    cfg.set_timing_cache(tc, ignore_mismatch=False)
+
+    profile = builder.create_optimization_profile()
+    profile.set_shape("wav", (1, T), (opt_batch, T), (max_batch, T))
+    cfg.add_optimization_profile(profile)
+
+    plan = builder.build_serialized_network(network, cfg)
+    if plan is None:
+        raise RuntimeError("engine build returned None")
+    Path(engine_path).write_bytes(bytes(plan))
+    Path(timing_cache).write_bytes(bytes(tc.serialize()))
+    return engine_path
+
+
+# ----------------------------------------------------------------------------
+# runtime
+# ----------------------------------------------------------------------------
+class _TRTCallable:
+    """Mimics net(xt): runs the engine on a [B,T] float32 CUDA tensor."""
+
+    def __init__(self, engine_path, n_classes, L, i):
+        import tensorrt as trt
+
+        self.n_classes = n_classes
+        self.rebuild = _rebuild(L, i)
+        logger = trt.Logger(trt.Logger.WARNING)
+        self.runtime = trt.Runtime(logger)
+        self.engine = self.runtime.deserialize_cuda_engine(Path(engine_path).read_bytes())
+        self.ctx = self.engine.create_execution_context()
+        if self.ctx is None:
+            raise RuntimeError(
+                "could not create execution context (likely OOM reserving max-profile "
+                "memory) — lower MAX_BATCH_CAP")
+        # resolve I/O tensor names
+        self.in_name, self.out_name = "wav", "logits"
+
+    def __call__(self, x, *a, **k):
+        x = x.to("cuda", torch.float32).reshape(x.shape[0], -1).contiguous()  # (B,1,T)->(B,T)
+        B = x.shape[0]
+        self.ctx.set_input_shape(self.in_name, tuple(x.shape))
+        out = torch.empty((B, self.n_classes), device="cuda", dtype=torch.float32)
+        self.ctx.set_tensor_address(self.in_name, x.data_ptr())
+        self.ctx.set_tensor_address(self.out_name, out.data_ptr())
+        stream = torch.cuda.current_stream().cuda_stream
+        self.ctx.execute_async_v3(stream)
+        torch.cuda.current_stream().synchronize()
+        return self.rebuild(out)
+
+
+# ----------------------------------------------------------------------------
+# parity + sweep
+# ----------------------------------------------------------------------------
+def _chunked_scores(model, audios, srs, chunk):
+    out = []
+    for k in range(0, len(audios), chunk):
+        out.extend(model.score_batch(audios[k:k + chunk], srs[k:k + chunk]))
+    return np.asarray(out, dtype=np.float64)
+
+
+def _spearman(a, b):
+    if len(a) < 2:
+        return 1.0
+    ra = np.argsort(np.argsort(a)).astype(np.float64)
+    rb = np.argsort(np.argsort(b)).astype(np.float64)
+    return float(np.corrcoef(ra, rb)[0, 1])
+
+
+def parity(model, trt_call, n=64, chunk=PARITY_CHUNK):
+    data = real_audio(n)
+    audios = [a for a, _ in data]
+    srs = [s for _, s in data]
+    torch_net = model.net
+    py = _chunked_scores(model, audios, srs, chunk)
+    model.net = trt_call
+    tr = _chunked_scores(model, audios, srs, chunk)
+    model.net = torch_net
+    mad = float(np.max(np.abs(py - tr)))
+    pear = float(np.corrcoef(py, tr)[0, 1]) if len(py) > 1 else 1.0
+    spear = _spearman(py, tr)
+    return {"n": len(py), "max_abs_score_diff": mad, "pearson": pear,
+            "spearman": spear}
+
+
+def sweep(model, trt_call,
+          batches=(1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128), iters=20):
+    a, sr = real_audio(1)[0]
+    model.net = trt_call
+    res = {}
+    for B in batches:
+        if B > MAX_BATCH_CAP:
+            break
+        ab, sb = [a] * B, [sr] * B
+        try:
+            for _ in range(3):
+                model.score_batch(ab, sb)                    # warmup
+            torch.cuda.synchronize()
+            t0 = time.time()
+            for _ in range(iters):
+                model.score_batch(ab, sb)
+            torch.cuda.synchronize()
+            dt = time.time() - t0
+            res[B] = B * iters / dt
+        except RuntimeError as e:
+            if "out of memory" in str(e).lower():
+                torch.cuda.empty_cache()
+                break
+            raise
+    best = max(res, key=res.get)
+    return best, res
+
+
+# ----------------------------------------------------------------------------
+# drop-in inference class
+# ----------------------------------------------------------------------------
+class DPDAGATCMTRT(_OrigClass):
+    """Drop-in: original preprocessing/score_batch; net replaced by the TRT engine."""
+
+    def load(self):
+        self.device = "cuda"
+        side = json.loads((HERE / f"trt_{SLUG}.json").read_text())[gpu_slug()]
+        eng = HERE / side["engine"]
+        self.net = _TRTCallable(str(eng), side["n_classes"], side["L"], side["i"])
+        self.batch_size = side["best_batch"]
+
+
+# ----------------------------------------------------------------------------
+# CLI
+# ----------------------------------------------------------------------------
+def _do_export():
+    gpu = gpu_slug()
+    side_path = HERE / f"trt_{SLUG}.json"
+    tc = HERE / f".trt_timing_{gpu}.cache"
+    m = load_model()
+    T, L, i, n_classes = analyze(m)
+    print(f"[analyze] T={T} n_classes={n_classes} L={L} i={i}")
+    onnx_path = HERE / f"{SLUG}.onnx"
+    export_onnx(m, T, L, i, onnx_path)
+    print(f"[onnx] wrote {onnx_path.name}")
+
+    # PyTorch reference scores while the model is on GPU, then free it so the
+    # engine build + TRT inference never co-reside with the model (giant >2 GB
+    # models would otherwise OOM the 16 GB card).
+    pdata = real_audio(64)
+    paud, psr = [a for a, _ in pdata], [s for _, s in pdata]
+    py = _chunked_scores(m, paud, psr, PARITY_CHUNK)
+    m.net.to("cpu")
+    torch.cuda.empty_cache()
+
+    opt_batch = min(32, MAX_BATCH_CAP)
+    if FORCE_FP16:
+        precisions = ["fp16"]
+    elif FORCE_FP32:
+        precisions = ["fp32"]
+    else:
+        precisions = ["fp16", "fp32"]
+    chosen = None
+    last_err = None
+    for prec in precisions:
+        eng = HERE / f"engine_{gpu}_{prec}_b1-{opt_batch}-{MAX_BATCH_CAP}.plan"
+        try:
+            t0 = time.time()
+            build_engine(str(onnx_path), T, prec, MAX_BATCH_CAP, opt_batch, str(eng), str(tc))
+            bt = time.time() - t0
+            trt_call = _TRTCallable(str(eng), n_classes, L, i)
+            m.net = trt_call
+            tr = _chunked_scores(m, paud, psr, PARITY_CHUNK)
+            p = {"n": len(py),
+                 "max_abs_score_diff": float(np.max(np.abs(py - tr))),
+                 "pearson": float(np.corrcoef(py, tr)[0, 1]) if len(py) > 1 else 1.0,
+                 "spearman": _spearman(py, tr)}
+        except Exception as e:  # noqa: BLE001 — try the next precision (e.g. FP16 layer not buildable)
+            last_err = e
+            print(f"[{prec}] FAILED: {type(e).__name__}: {e}")
+            continue
+        print(f"[{prec}] build={bt:.1f}s parity={p}")
+        chosen = (prec, eng, p, trt_call)
+        if prec == "fp16" and p["spearman"] >= PARITY_SPEARMAN_TOL:
+            break
+
+    if chosen is None:
+        raise RuntimeError(f"all precisions failed to build; last error: {last_err}")
+    prec, eng, p, trt_call = chosen
+    if p["spearman"] < PARITY_FLOOR:
+        raise RuntimeError(
+            f"parity too low (spearman={p['spearman']:.4f} < {PARITY_FLOOR}): "
+            f"engine output does not match PyTorch — not accepting")
+    m.net = trt_call
+    best, table = sweep(m, trt_call)
+    side = json.loads(side_path.read_text()) if side_path.exists() else {}
+    side[gpu] = {
+        "precision": prec, "engine": eng.name, "window_samples": T,
+        "n_classes": n_classes, "L": L, "i": i, "best_batch": best,
+        "throughput_utt_s": {str(k): round(v, 2) for k, v in table.items()},
+        "parity": p, "trt_version": __import__("tensorrt").__version__,
+    }
+    side_path.write_text(json.dumps(side, indent=2, default=str))
+    print(f"[done] {SLUG}: prec={prec} best_batch={best} "
+          f"utt/s={table[best]:.1f} parity_mad={p['max_abs_score_diff']:.2e}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("cmd", choices=["export", "sweep", "parity", "score"])
+    ap.add_argument("audio", nargs="?")
+    args = ap.parse_args()
+    gpu = gpu_slug()
+    side_path = HERE / f"trt_{SLUG}.json"
+
+    if args.cmd == "export":
+        _do_export()
+    elif args.cmd in ("sweep", "parity"):
+        m = load_model()
+        side = json.loads(side_path.read_text())[gpu]
+        eng = HERE / side["engine"]
+        trt_call = _TRTCallable(str(eng), side["n_classes"], side["L"], side["i"])
+        if args.cmd == "parity":
+            print(parity(m, trt_call))
+        else:
+            best, table = sweep(m, trt_call)
+            full = json.loads(side_path.read_text())
+            full[gpu]["best_batch"] = best
+            full[gpu]["throughput_utt_s"] = {str(k): round(v, 2) for k, v in table.items()}
+            side_path.write_text(json.dumps(full, indent=2, default=str))
+            print(f"best_batch={best} utt/s={table[best]:.1f}")
+    elif args.cmd == "score":
+        import soundfile as sf
+        a, sr = sf.read(args.audio, dtype="float32")
+        if a.ndim > 1:
+            a = a.mean(1)
+        m = DPDAGATCMTRT()
+        m.load()
+        print(m.score_batch([a.astype(np.float32)], [sr])[0])
+
+
+if __name__ == "__main__":
+    main()