diff --git a/.gitignore b/.gitignore index 1c4c122..7902802 100644 --- a/.gitignore +++ b/.gitignore @@ -47,6 +47,7 @@ opendpd_env/ dpd_out/ log/ save/ +benchmark/results/reproduced/ # Allow sample data but not large datasets datasets/**/data/*.csv @@ -94,4 +95,4 @@ steps/__pycache__/ PACKAGE_INSTALLATION.md PyPI.md -adc/* \ No newline at end of file +adc/* diff --git a/arguments.py b/arguments.py index 608e3ce..3dc2fe2 100644 --- a/arguments.py +++ b/arguments.py @@ -31,22 +31,32 @@ def get_arguments(): parser.add_argument('--seed', default=0, type=int, help='Global random number seed.') parser.add_argument('--loss_type', default='l2', choices=['l1', 'l2'], help='Type of loss function.') parser.add_argument('--opt_type', default='adamw', choices=['sgd', 'adam', 'adamw', 'adabound', 'rmsprop'], help='Type of optimizer.') - parser.add_argument('--batch_size', default=256, type=int, help='Batch size for training.') - parser.add_argument('--batch_size_eval', default=256, type=int, help='Batch size for evaluation.') - parser.add_argument('--n_epochs', default=100, type=int, help='Number of epochs to train for.') - parser.add_argument('--lr_schedule', default=0, type=int, help='Whether enable learning rate scheduling') - parser.add_argument('--lr', default=5e-4, type=float, help='Learning rate') - parser.add_argument('--lr_end', default=1e-4, type=float, help='Learning rate') - parser.add_argument('--decay_factor', default=0.1, type=float, help='Learning rate') - parser.add_argument('--patience', default=10, type=float, help='Learning rate') + parser.add_argument('--batch_size', default=64, type=int, help='Batch size for training.') + parser.add_argument('--batch_size_eval', default=64, type=int, help='Batch size for evaluation.') + parser.add_argument('--n_epochs', default=300, type=int, help='Number of epochs to train for.') + parser.add_argument('--lr_schedule', default=1, type=int, + help='Whether to enable ReduceLROnPlateau learning-rate scheduling.') + parser.add_argument('--lr', default=5e-3, type=float, help='Initial learning rate.') + parser.add_argument('--lr_end', default=5e-5, type=float, help='Minimum learning rate.') + parser.add_argument('--decay_factor', default=0.5, type=float, + help='Learning-rate reduction factor.') + parser.add_argument('--patience', default=5, type=int, + help='Scheduler patience in epochs.') parser.add_argument('--grad_clip_val', default=200, type=float, help='Gradient clipping.') + parser.add_argument( + '--cuda_graph_training', action='store_true', default=False, + help=( + 'Opt in to guarded whole-step CUDA-graph replay for supported ' + 'TRes-DeltaGRU DPD training.' + ), + ) # GMP Hyperparameters parser.add_argument('--K', default=5, type=int, help='Degree of GMP model') parser.add_argument('--gmp_memory_length', default=11, type=int, help='Memory length of GMP model') # Power Amplifier Model Settings parser.add_argument('--PA_backbone', default='gru', choices=['gmp','deltagru', 'deltajanet', 'janet', 'fcn', 'gru', 'dgru', 'qgru', 'qgru_amp1', 'lstm', 'vdlstm', - 'rvtdcnn', 'mamba', 'tcn', 'pntdnn', 'pdgru', 'pgjanet', 'dvrjanet', 'bojanet', 'pnjanet', 'apnrnn', 'djanet', 'mcldnn'], + 'rvtdcnn', 'mamba', 'tcn', 'tres_deltagru', 'tres_gru', 'pntdnn', 'pdgru', 'pgjanet', 'dvrjanet', 'bojanet', 'pnjanet', 'apnrnn', 'djanet', 'mcldnn'], help='Modeling PA Recurrent layer type') parser.add_argument('--PA_hidden_size', default=23, type=int, help='Hidden size of PA backbone') @@ -55,7 +65,7 @@ def get_arguments(): # Digital Predistortion Model Settings parser.add_argument('--DPD_backbone', default='gru', choices=['gmp', 'deltagru', 'deltajanet', 'janet', 'snn', 'fcn', 'gru', 'dgru', 'qgru', 'qgru_amp1', 'lstm', 'vdlstm', - 'rvtdcnn', 'tres_deltagru', 'tcn', 'pntdnn', 'pdgru', 'pgjanet', 'dvrjanet', 'bojanet', 'pnjanet', 'djanet', 'mcldnn'], + 'rvtdcnn', 'tres_deltagru', 'tres_gru', 'tcn', 'pntdnn', 'pdgru', 'pgjanet', 'dvrjanet', 'bojanet', 'pnjanet', 'djanet', 'mcldnn'], help='DPD model Recurrent layer type') parser.add_argument('--DPD_hidden_size', default=15, type=int, help='Hidden size of DPD backbone.') parser.add_argument('--DPD_num_layers', default=1, type=int, help='Number of layers of the DPD backbone.') @@ -76,6 +86,8 @@ def get_arguments(): help='Threshold for input deltas') parser.add_argument('--thh', type=float, default=0.0, help='Threshold for hidden state deltas') + parser.add_argument('--collect_delta_stats', action='store_true', default=False, + help='Collect temporal delta sparsity diagnostics during training (adds overhead).') # Optionally, you might want to add DVR-specific arguments parser.add_argument('--num_dvr_units', default=3, type=int, @@ -94,4 +106,4 @@ def get_arguments(): parser.add_argument('--gif_duration', default=10.0, type=float, help='Duration of GIF animations in seconds (default: 10.0).') - return parser.parse_args() \ No newline at end of file + return parser.parse_args() diff --git a/backbones/cuda_graph_frozen_dgru.py b/backbones/cuda_graph_frozen_dgru.py new file mode 100644 index 0000000..51c2dbb --- /dev/null +++ b/backbones/cuda_graph_frozen_dgru.py @@ -0,0 +1,317 @@ +"""Exact CUDA-graph replay for a frozen DGRU used during DPD training. + +The optimization records the existing cuDNN-backed DGRU forward and its input +backward as two CUDA graphs. It neither replaces the recurrence nor changes +the module tree. Unsupported or unsafe calls return ``None`` so the caller can +execute its ordinary eager implementation. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import os +import threading +import weakref + +import torch + + +_DISABLE_ENV = "OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU" +_TRUE_VALUES = {"1", "true", "yes", "on"} +_MANAGERS = weakref.WeakKeyDictionary() +_MANAGERS_LOCK = threading.Lock() +_CAPTURE_FAILED = object() +_LOCAL_DISABLE_DEPTH = contextvars.ContextVar( + "opendpd_frozen_dgru_graph_disable_depth", default=0 +) + + +def _disabled() -> bool: + return ( + _LOCAL_DISABLE_DEPTH.get() > 0 + or os.getenv(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES + ) + + +@contextlib.contextmanager +def disable_cuda_graph_frozen_dgru(): + """Temporarily bypass the inner graph without changing process state.""" + + token = _LOCAL_DISABLE_DEPTH.set(_LOCAL_DISABLE_DEPTH.get() + 1) + try: + yield + finally: + _LOCAL_DISABLE_DEPTH.reset(token) + + +def _cuda_autocast_enabled() -> bool: + try: + return torch.is_autocast_enabled("cuda") + except TypeError: # pragma: no cover - compatibility with older PyTorch. + return torch.is_autocast_enabled() + + +def _can_replay(module, input: torch.Tensor, h_0: torch.Tensor | None) -> bool: + if ( + _disabled() + or not module.training + or not torch.is_grad_enabled() + or not input.requires_grad + or input.ndim != 3 + or not input.is_cuda + or input.dtype != torch.float32 + or not input.is_contiguous() + or torch.version.cuda is None + or getattr(torch.version, "hip", None) is not None + or _cuda_autocast_enabled() + or not hasattr(torch.cuda, "CUDAGraph") + ): + return False + try: + if torch.cuda.is_current_stream_capturing(): + return False + except RuntimeError: + return False + + parameters = tuple(module.parameters()) + if not parameters or any( + parameter.requires_grad + or parameter.device != input.device + or parameter.dtype != input.dtype + for parameter in parameters + ): + return False + if h_0 is not None and ( + h_0.requires_grad + or h_0.device != input.device + or h_0.dtype != input.dtype + or not h_0.is_contiguous() + ): + return False + return True + + +def _cache_key(module, input: torch.Tensor, h_0: torch.Tensor | None): + parameter_storage = tuple( + parameter.data_ptr() for parameter in module.parameters() + ) + h0_signature = ( + None if h_0 is None else (tuple(h_0.shape), tuple(h_0.stride())) + ) + return ( + input.device.index, + input.dtype, + tuple(input.shape), + tuple(input.stride()), + h0_signature, + parameter_storage, + ) + + +class _ReplayCache: + def __init__(self): + self.entries = {} + self.lock = threading.Lock() + + def get_or_create(self, module, input, h_0): + key = _cache_key(module, input, h_0) + entry = self.entries.get(key) + if entry is _CAPTURE_FAILED: + return None + if entry is not None: + return entry + + # Capturing is expensive and unsafe to duplicate. A concurrent first + # call simply uses eager execution and can use the cache next time. + if not self.lock.acquire(blocking=False): + return None + try: + entry = self.entries.get(key) + if entry is _CAPTURE_FAILED: + return None + if entry is None: + try: + entry = _FrozenDGRUReplay(module, input, h_0) + except (RuntimeError, torch.cuda.OutOfMemoryError): + self.entries[key] = _CAPTURE_FAILED + return None + self.entries[key] = entry + return entry + finally: + self.lock.release() + + +class _FrozenDGRUReplay: + def __init__(self, module, example_input, example_h0): + self.lock = threading.Lock() + self.busy = False + self.failed = False + self.generation = 0 + self.pending_event = None + self.pending_stream = None + + with torch.cuda.device(example_input.device), torch.enable_grad(): + # Warm cuDNN and autograd on a side stream before capture. Using + # that same stream for both captures keeps retained autograd nodes + # from acquiring stale cross-stream dependencies. + self.capture_stream = torch.cuda.Stream( + device=example_input.device + ) + self.capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.capture_stream): + self.static_input = ( + example_input.detach().clone().requires_grad_() + ) + self.static_h0 = ( + None if example_h0 is None else example_h0.detach().clone() + ) + for _ in range(3): + warm_output = module._forward_eager( + self.static_input, self.static_h0 + ) + warm_gradient = torch.ones_like(warm_output) + torch.autograd.grad( + warm_output, + self.static_input, + grad_outputs=warm_gradient, + ) + del warm_output, warm_gradient + torch.cuda.current_stream().wait_stream(self.capture_stream) + torch.cuda.synchronize(example_input.device) + + self.forward_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + self.forward_graph, stream=self.capture_stream + ): + self.static_output = module._forward_eager( + self.static_input, self.static_h0 + ) + + self.static_grad_output = torch.empty_like(self.static_output) + self.backward_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + self.backward_graph, stream=self.capture_stream + ): + (self.static_grad_input,) = torch.autograd.grad( + self.static_output, + self.static_input, + grad_outputs=self.static_grad_output, + ) + + def reserve(self) -> int | None: + with self.lock: + if self.failed or self.busy: + return None + if self.pending_event is not None: + current_stream = torch.cuda.current_stream( + self.static_input.device + ) + same_stream = current_stream.cuda_stream == self.pending_stream + if not same_stream and not self.pending_event.query(): + return None + self.pending_event = None + self.pending_stream = None + self.busy = True + self.generation += 1 + return self.generation + + def release(self, generation: int, *, failed: bool = False) -> None: + with self.lock: + if generation != self.generation: + return + self.failed = self.failed or failed + if torch.cuda.is_available(): + event = torch.cuda.Event() + stream = torch.cuda.current_stream(self.static_input.device) + event.record(stream) + self.pending_event = event + self.pending_stream = stream.cuda_stream + self.busy = False + + def forward(self, input, h_0, generation): + if generation != self.generation or not self.busy: + raise RuntimeError("stale frozen-DGRU CUDA-graph reservation") + self.static_input.copy_(input) + if self.static_h0 is not None: + self.static_h0.copy_(h_0) + self.forward_graph.replay() + return self.static_output.clone() + + def backward(self, grad_output, generation): + if generation != self.generation or not self.busy: + raise RuntimeError( + "frozen-DGRU CUDA graph cannot be replayed twice" + ) + self.static_grad_output.copy_(grad_output) + self.backward_graph.replay() + return self.static_grad_input.clone() + + +class _FrozenDGRUReplayFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, h_0, replay, generation): + ctx.replay = replay + ctx.generation = generation + ctx.used = False + try: + return replay.forward(input, h_0, generation) + except Exception: + replay.release(generation, failed=True) + raise + + @staticmethod + def backward(ctx, grad_output): + replay = ctx.replay + generation = ctx.generation + if ctx.used: + raise RuntimeError( + "frozen-DGRU CUDA graph supports one backward pass" + ) + ctx.used = True + if torch.is_grad_enabled(): + replay.release(generation) + raise RuntimeError( + "CUDA-graph frozen DGRU supports first-order training only; " + f"set {_DISABLE_ENV}=1 for higher-order autograd" + ) + try: + grad_input = replay.backward(grad_output, generation) + except Exception: + replay.release(generation, failed=True) + raise + replay.release(generation) + return grad_input, None, None, None + + +def try_cuda_graph_frozen_dgru( + module, input: torch.Tensor, h_0: torch.Tensor | None +) -> torch.Tensor | None: + """Return a graph-replayed output, or ``None`` for eager fallback.""" + + if not _can_replay(module, input, h_0): + return None + with _MANAGERS_LOCK: + manager = _MANAGERS.get(module) + if manager is None: + manager = _ReplayCache() + _MANAGERS[module] = manager + replay = manager.get_or_create(module, input, h_0) + if replay is None: + return None + generation = replay.reserve() + if generation is None: + return None + try: + return _FrozenDGRUReplayFunction.apply(input, h_0, replay, generation) + except (RuntimeError, torch.cuda.OutOfMemoryError): + # A replay error disables this cache entry. Eager execution still has + # the original input and is safe unless CUDA itself is unrecoverable. + return None + + +def clear_cuda_graph_frozen_dgru_cache() -> None: + """Drop replay caches; primarily useful for focused tests.""" + + with _MANAGERS_LOCK: + _MANAGERS.clear() diff --git a/backbones/dgru.py b/backbones/dgru.py index a3c7dd2..284e67b 100644 --- a/backbones/dgru.py +++ b/backbones/dgru.py @@ -7,8 +7,8 @@ class DGRU(nn.Module): - def __init__(self, hidden_size, output_size, num_layers, bidirectional=False, batch_first=True, - bias=True): + def __init__(self, hidden_size, output_size, num_layers, + bidirectional=False, batch_first=True, bias=True): super(DGRU, self).__init__() self.hidden_size = hidden_size self.input_size = 6 @@ -39,10 +39,14 @@ def reset_parameters(self): nn.init.constant_(param, 0) if 'weight' in name: for i in range(0, num_gates): - nn.init.orthogonal_(param[i * self.hidden_size:(i + 1) * self.hidden_size, :]) + start = i * self.hidden_size + end = (i + 1) * self.hidden_size + nn.init.orthogonal_(param[start:end, :]) if 'weight_ih_l0' in name: for i in range(0, num_gates): - nn.init.xavier_uniform_(param[i * self.hidden_size:(i + 1) * self.hidden_size, :]) + start = i * self.hidden_size + end = (i + 1) * self.hidden_size + nn.init.xavier_uniform_(param[start:end, :]) for name, param in self.fc_out.named_parameters(): if 'weight' in name: @@ -56,7 +60,7 @@ def reset_parameters(self): if 'bias' in name: nn.init.constant_(param, 0) - def forward(self, x, h_0): + def _forward_eager(self, x, h_0): # Feature Extraction i_x = torch.unsqueeze(x[..., 0], dim=-1) q_x = torch.unsqueeze(x[..., 1], dim=-1) @@ -72,3 +76,15 @@ def forward(self, x, h_0): out = torch.cat((out, x), dim=-1) out = self.fc_out(out) return out + + def forward(self, x, h_0): + # A frozen PA needs gradients only with respect to its input. For the + # common fixed-shape CUDA training path, replay the exact eager/cuDNN + # kernels and their input backward to remove launch overhead. The + # helper returns None for every unsupported or unsafe call. + from backbones.cuda_graph_frozen_dgru import try_cuda_graph_frozen_dgru + + graph_output = try_cuda_graph_frozen_dgru(self, x, h_0) + if graph_output is not None: + return graph_output + return self._forward_eager(x, h_0) diff --git a/backbones/tres_deltagru.py b/backbones/tres_deltagru.py index 2755bb3..6699c62 100644 --- a/backbones/tres_deltagru.py +++ b/backbones/tres_deltagru.py @@ -1,12 +1,19 @@ +import math +import os + +import numpy as np import torch -from torch import Tensor import torch.nn as nn -import torch.nn.functional as F -import numpy as np -import csv -import os +from torch import Tensor + from quant.modules.ops import Add, Mul -from torch import nn + +try: + from backbones.triton_deltagru import can_use_triton_deltagru, triton_deltagru +except ImportError: # Triton is optional; the eager implementation is portable. + can_use_triton_deltagru = None + triton_deltagru = None + class TResDeltaGRU(nn.Module): def __init__(self, input_size, hidden_size, output_size, num_layers, thx=0, thh=0, @@ -47,10 +54,16 @@ def __init__(self, input_size, hidden_size, output_size, num_layers, thx=0, thh= bias=False), nn.Hardswish(), ) - self.set_debug(1) + # Statistics are diagnostic side effects, not part of the network. They + # used to add four reductions per recurrent timestep during every train + # step. Keep them opt-in and collect them inside the fused kernel when + # explicitly enabled. + self.set_debug(0) def set_debug(self, value): + value = int(value) setattr(self, "debug", value) + self.rnn.debug = value self.rnn.statistics = { "num_dx_zeros": 0, "num_dx_numel": 0, @@ -86,7 +99,7 @@ def reset_parameters(self): nn.init.constant_(param, 0) - def forward(self, x, h_0): + def forward(self, x, h_0=None): skip_x = self.tcn(x.transpose(1,2)).transpose(1,2) last_step = torch.roll(x, shifts=-1, dims=1) # Feature Extraction @@ -98,7 +111,11 @@ def forward(self, x, h_0): amp = torch.sqrt(amp2) amp3 = torch.pow(amp, 3) x = torch.cat((i_x, q_x, amp, amp3, i_last, q_last), dim=-1) - out = self.rnn(x, h_0) + # The old call passed ``h_0`` as x_p_0 while leaving the other recurrent + # states unset; process_inputs_first consequently discarded it and made + # all-zero states. Calling without it preserves that behaviour and + # avoids an unnecessary allocation/copy in CoreModel. + out = self.rnn(x) out = self.fc_out(out)+skip_x return out @@ -147,6 +164,7 @@ def __init__(self, self.x_p_length = max(self.input_size, self.hidden_size) self.batch_first = True self.debug = 1 + self.use_triton = os.environ.get("OPENDPD_DISABLE_TRITON_DELTAGRU", "0") != "1" # Statistics self.abs_sum_delta_hid = torch.zeros(1) self.sp_dx = 0 @@ -246,29 +264,23 @@ def layer_forward(self, input: Tensor, l: int, x_p_0: Tensor = None, h_0: Tensor input_size = input.size(-1) - th_x = torch.tensor(self.th_x, dtype=input.dtype) - th_h = torch.tensor(self.th_h, dtype=input.dtype) + th_x = self.th_x + th_h = self.th_h inputs = input.unbind(0) output = [] - reg = torch.zeros(1, dtype=input.dtype, device=input.device).squeeze() - x_p = x_p_0[:, :input_size] h = h_0 h_p = h_p_0 dm_nh = dm_nh_0 dm = dm_0 - l1_norm_delta_h = torch.zeros(1, dtype=input.dtype, device=input.device) - seq_len = len(inputs) for t in range(seq_len): x = inputs[t] delta_x, delta_h, delta_x_abs, delta_h_abs = self.compute_deltas(x, x_p, h, h_p, th_x, th_h) - reg += torch.sum(torch.abs(delta_h)) - if self.debug: zero_mask_delta_x = torch.as_tensor(delta_x == 0, dtype=x.dtype) zero_mask_delta_h = torch.as_tensor(delta_h == 0, dtype=x.dtype) @@ -279,8 +291,6 @@ def layer_forward(self, input: Tensor, l: int, x_p_0: Tensor = None, h_0: Tensor x_p, h_p = self.update_states(delta_x_abs, delta_h_abs, x, h, x_p, h_p, th_x, th_h) - l1_norm_delta_h += torch.sum(torch.abs(delta_h)) - dm, dm_r, dm_z, dm_n, dm_nh = self.compute_gates(delta_x, delta_h, dm, dm_nh) gate_r = self.sigmoid(dm_r) @@ -295,10 +305,107 @@ def layer_forward(self, input: Tensor, l: int, x_p_0: Tensor = None, h_0: Tensor def forward(self, input: Tensor, x_p_0: Tensor = None, h_0: Tensor = None, h_p_0: Tensor = None, dm_nh_0: Tensor = None, dm_0: Tensor = None): + states = (x_p_0, h_0, h_p_0, dm_nh_0, dm_0) + if self._can_use_triton(input) and self._states_can_use_triton(input, states): + batch_size = input.size(0) + if x_p_0 is None or h_0 is None or h_p_0 is None or dm_nh_0 is None or dm_0 is None: + x_p_0 = input.new_zeros((self.num_layers, batch_size, self.x_p_length)) + h_0 = input.new_zeros((self.num_layers, batch_size, self.hidden_size)) + h_p_0 = input.new_zeros((self.num_layers, batch_size, self.hidden_size)) + dm_nh_0 = input.new_zeros((self.num_layers, batch_size, self.hidden_size)) + dm_0 = input.new_zeros((self.num_layers, batch_size, self.weight_ih_height)) + + collect_stats = bool(self.debug) + stats = torch.zeros(2, dtype=torch.int64, device=input.device) if collect_stats else torch.empty( + 2, dtype=torch.int64, device=input.device + ) + output = triton_deltagru( + input.contiguous(), + x_p_0[0, :, :self.input_size].contiguous(), + h_0[0].contiguous(), + h_p_0[0].contiguous(), + dm_nh_0[0].contiguous(), + dm_0[0].contiguous(), + self.x2h.weight, + self.h2h.weight, + stats, + self.th_x, + self.th_h, + collect_stats, + ) + if collect_stats: + # Match the eager diagnostic counters, which use input dtype. + self.statistics["num_dx_zeros"] += stats[0].to(input.dtype) + self.statistics["num_dh_zeros"] += stats[1].to(input.dtype) + self.statistics["num_dx_numel"] += input.numel() + self.statistics["num_dh_numel"] += batch_size * input.size(1) * self.hidden_size + return output + x, x_p_0, h_0, h_p_0, dm_nh_0, dm_0 = self.process_inputs_first(input, x_p_0, h_0, h_p_0, dm_nh_0, dm_0) for l in range(self.num_layers): x = self.layer_forward(x, l, x_p_0[l], h_0[l], h_p_0[l], dm_nh_0[l], dm_0[l]) x = x.transpose(0, 1) - return x \ No newline at end of file + return x + + def _can_use_triton(self, input: Tensor) -> bool: + """The fused path is for the dense-float, one-layer CUDA cell. + + Quantization-aware training replaces these exact module types with + fake-quantized subclasses/operations. Reassociating those quantization + boundaries would change the model, so QAT deliberately stays eager. + """ + + return ( + self.use_triton + and self.num_layers == 1 + and can_use_triton_deltagru is not None + and not torch.is_autocast_enabled() + and not math.isnan(self.th_x) + and not math.isnan(self.th_h) + and input.ndim == 3 + and input.size(-1) == self.input_size + and type(self.x2h) is nn.Linear + and type(self.h2h) is nn.Linear + and self.x2h.bias is None + and self.h2h.bias is None + and self.x2h.weight.shape == (self.weight_ih_height, self.input_size) + and self.h2h.weight.shape == (self.weight_ih_height, self.hidden_size) + and self.x2h.weight.device == input.device + and self.h2h.weight.device == input.device + and self.x2h.weight.dtype == input.dtype + and self.h2h.weight.dtype == input.dtype + and self.x2h.weight.is_contiguous() + and self.h2h.weight.is_contiguous() + and type(self.sigmoid) is nn.Sigmoid + and type(self.tanh) is nn.Tanh + and type(self.add) is Add + and type(self.mul) is Mul + and can_use_triton_deltagru(input, self.hidden_size) + ) + + def _states_can_use_triton(self, input: Tensor, states: tuple) -> bool: + # The eager path discards every supplied state when any one is missing. + # The fused path does the same, so that case is always safe. + if any(state is None for state in states): + return True + + x_p_0, h_0, h_p_0, dm_nh_0, dm_0 = states + if not all(isinstance(state, Tensor) for state in states): + return False + batch_size = input.size(0) + common = (self.num_layers, batch_size) + expected_shapes = ( + common + (self.x_p_length,), + common + (self.hidden_size,), + common + (self.hidden_size,), + common + (self.hidden_size,), + common + (self.weight_ih_height,), + ) + return all( + tuple(state.shape) == shape + and state.device == input.device + and state.dtype == input.dtype + for state, shape in zip(states, expected_shapes) + ) diff --git a/backbones/tres_gru.py b/backbones/tres_gru.py new file mode 100644 index 0000000..7c5c789 --- /dev/null +++ b/backbones/tres_gru.py @@ -0,0 +1,85 @@ +import torch +import torch.nn as nn + + +class TResGRU(nn.Module): + """TRes-DeltaGRU with the DeltaGRU cell replaced by a dense nn.GRU. + + Same 6-feature extraction and TCN residual skip path as TResDeltaGRU; + the recurrent core is a standard cuDNN GRU. Bias-free like the DeltaGRU + layer's x2h/h2h projections, so parameter counts match at equal hidden + size (e.g. H=10 -> 524 params). + """ + + def __init__(self, input_size, hidden_size, output_size, num_layers, bias=True): + super(TResGRU, self).__init__() + self.hidden_size = hidden_size + self.input_size = 6 + self.output_size = output_size + self.num_layers = num_layers + self.bias = bias + + self.rnn = nn.GRU(input_size=self.input_size, + hidden_size=self.hidden_size, + num_layers=num_layers, + batch_first=True, + bias=False) + self.fc_out = nn.Linear(in_features=self.hidden_size, + out_features=self.output_size, + bias=False) + self.tcn = nn.Sequential( + nn.Conv1d(in_channels=2, + out_channels=3, + kernel_size=3, + padding=16, + stride=1, + dilation=16, + bias=False), + nn.Hardswish(), + nn.Conv1d(in_channels=3, + out_channels=2, + kernel_size=1, + padding=0, + stride=1, + dilation=1, + bias=False), + nn.Hardswish(), + ) + + def reset_parameters(self): + for name, param in self.tcn.named_parameters(): + if 'weight' in name: + nn.init.xavier_uniform_(param) + if 'bias' in name: + nn.init.constant_(param, 0) + for name, param in self.rnn.named_parameters(): + num_gates = int(param.shape[0] / self.hidden_size) + if 'bias' in name: + nn.init.constant_(param, 0) + if 'weight_hh' in name: + for i in range(0, num_gates): + nn.init.orthogonal_(param[i * self.hidden_size:(i + 1) * self.hidden_size, :]) + if 'weight_ih' in name: + for i in range(0, num_gates): + nn.init.xavier_uniform_(param[i * self.hidden_size:(i + 1) * self.hidden_size, :]) + for name, param in self.fc_out.named_parameters(): + if 'weight' in name: + nn.init.xavier_uniform_(param) + if 'bias' in name: + nn.init.constant_(param, 0) + + def forward(self, x, h_0=None): + skip_x = self.tcn(x.transpose(1, 2)).transpose(1, 2) + last_step = torch.roll(x, shifts=-1, dims=1) + # Feature Extraction + i_x = x[..., 0].unsqueeze(-1) + q_x = x[..., 1].unsqueeze(-1) + i_last = last_step[..., 0].unsqueeze(-1) + q_last = last_step[..., 1].unsqueeze(-1) + amp2 = torch.pow(i_x, 2) + torch.pow(q_x, 2) + amp = torch.sqrt(amp2) + amp3 = torch.pow(amp, 3) + feat = torch.cat((i_x, q_x, amp, amp3, i_last, q_last), dim=-1) + out, _ = self.rnn(feat, h_0) + out = self.fc_out(out) + skip_x + return out diff --git a/backbones/triton_deltagru.py b/backbones/triton_deltagru.py new file mode 100644 index 0000000..861a0f1 --- /dev/null +++ b/backbones/triton_deltagru.py @@ -0,0 +1,632 @@ +"""Fused CUDA implementation of the DeltaGRU recurrence. + +The public helper in this module is intentionally small and optional. Importing +``backbones.tres_deltagru`` must continue to work on machines without Triton; +callers should check :func:`can_use_triton_deltagru` before dispatching here. + +The kernels preserve the eager recurrence, including its threshold comparisons +and separate input/hidden accumulator states. The backward kernel performs the +reverse recurrence and leaves the two weight-gradient reductions to cuBLAS. +That avoids a highly contended atomic reduction over batch and time. +""" + +from __future__ import annotations + +import math + +import torch + +try: + import triton + import triton.language as tl + from triton.language.extra import libdevice +except (ImportError, OSError): # pragma: no cover - CPU-only installations. + triton = None + tl = None + libdevice = None + + +def can_use_triton_deltagru(input: torch.Tensor, hidden_size: int) -> bool: + """Return whether the fused kernel supports this tensor configuration.""" + + return ( + triton is not None + and input.ndim == 3 + and input.shape[0] > 0 + and input.shape[1] > 0 + and input.is_cuda + and input.is_contiguous() + and torch.version.cuda is not None + and getattr(torch.version, "hip", None) is None + and input.dtype == torch.float32 + and 0 < input.shape[-1] <= 16 + and 0 < hidden_size <= 32 + ) + + +if triton is not None: + + @triton.jit + def _deltagru_forward_kernel( + x_ptr, + x_p0_ptr, + h0_ptr, + h_p0_ptr, + dm_nh0_ptr, + dm0_ptr, + weight_x_ptr, + weight_h_ptr, + out_ptr, + delta_x_ptr, + delta_h_ptr, + gate_r_ptr, + gate_z_ptr, + gate_n_ptr, + dm_nh_hist_ptr, + mask_x_ptr, + mask_h_ptr, + stats_ptr, + seq_len, + input_size: tl.constexpr, + hidden_size: tl.constexpr, + th_x: tl.constexpr, + th_h: tl.constexpr, + collect_stats: tl.constexpr, + save_intermediates: tl.constexpr, + block_i: tl.constexpr, + block_h: tl.constexpr, + ): + batch = tl.program_id(0) + offs_i = tl.arange(0, block_i) + offs_h = tl.arange(0, block_h) + mask_i = offs_i < input_size + mask_h = offs_h < hidden_size + + # Every program owns one batch row, so all recurrent state stays local. + x_p = tl.load(x_p0_ptr + batch * input_size + offs_i, mask=mask_i, other=0.0) + h = tl.load(h0_ptr + batch * hidden_size + offs_h, mask=mask_h, other=0.0) + h_p = tl.load(h_p0_ptr + batch * hidden_size + offs_h, mask=mask_h, other=0.0) + dm_r = tl.load(dm0_ptr + batch * (3 * hidden_size) + offs_h, mask=mask_h, other=0.0) + dm_z = tl.load( + dm0_ptr + batch * (3 * hidden_size) + hidden_size + offs_h, + mask=mask_h, + other=0.0, + ) + dm_n = tl.load( + dm0_ptr + batch * (3 * hidden_size) + 2 * hidden_size + offs_h, + mask=mask_h, + other=0.0, + ) + dm_nh = tl.load(dm_nh0_ptr + batch * hidden_size + offs_h, mask=mask_h, other=0.0) + + # The weights are loop-invariant. Loading them before the scan also + # gives the compiler a chance to keep this tiny model in registers. + wx_offsets = offs_h[:, None] * input_size + offs_i[None, :] + wh_offsets = offs_h[:, None] * hidden_size + offs_h[None, :] + wx_r = tl.load(weight_x_ptr + wx_offsets, mask=mask_h[:, None] & mask_i[None, :], other=0.0) + wx_z = tl.load( + weight_x_ptr + hidden_size * input_size + wx_offsets, + mask=mask_h[:, None] & mask_i[None, :], + other=0.0, + ) + wx_n = tl.load( + weight_x_ptr + 2 * hidden_size * input_size + wx_offsets, + mask=mask_h[:, None] & mask_i[None, :], + other=0.0, + ) + wh_r = tl.load(weight_h_ptr + wh_offsets, mask=mask_h[:, None] & mask_h[None, :], other=0.0) + wh_z = tl.load( + weight_h_ptr + hidden_size * hidden_size + wh_offsets, + mask=mask_h[:, None] & mask_h[None, :], + other=0.0, + ) + wh_n = tl.load( + weight_h_ptr + 2 * hidden_size * hidden_size + wh_offsets, + mask=mask_h[:, None] & mask_h[None, :], + other=0.0, + ) + + zeros_x = tl.zeros((), tl.int64) + zeros_h = tl.zeros((), tl.int64) + + for time in tl.range(0, seq_len, loop_unroll_factor=1): + x_offsets = (batch * seq_len + time) * input_size + offs_i + xt = tl.load(x_ptr + x_offsets, mask=mask_i, other=0.0) + + raw_delta_x = xt - x_p + raw_delta_h = h - h_p + abs_delta_x = tl.abs(raw_delta_x) + abs_delta_h = tl.abs(raw_delta_h) + + # masked_fill(abs(delta) < threshold, 0) and where(abs(delta) >= + # threshold, current, previous) are deliberately kept separate: + # their behaviour differs for NaNs in the eager implementation. + keep_delta_x = ~(abs_delta_x < th_x) + keep_delta_h = ~(abs_delta_h < th_h) + update_x = abs_delta_x >= th_x + update_h = abs_delta_h >= th_h + dx = tl.where(keep_delta_x, raw_delta_x, 0.0) + dh = tl.where(keep_delta_h, raw_delta_h, 0.0) + x_p = tl.where(update_x, xt, x_p) + h_p = tl.where(update_h, h, h_p) + + input_mac_r = tl.sum(wx_r * dx[None, :], axis=1) + dm_r + input_mac_z = tl.sum(wx_z * dx[None, :], axis=1) + dm_z + input_mac_n = tl.sum(wx_n * dx[None, :], axis=1) + dm_n + hidden_mac_r = tl.sum(wh_r * dh[None, :], axis=1) + hidden_mac_z = tl.sum(wh_z * dh[None, :], axis=1) + hidden_mac_n = tl.sum(wh_n * dh[None, :], axis=1) + + dm_r = input_mac_r + hidden_mac_r + dm_z = input_mac_z + hidden_mac_z + dm_n = input_mac_n + dm_nh = hidden_mac_n + dm_nh + + r = tl.sigmoid(dm_r) + z = tl.sigmoid(dm_z) + n = libdevice.tanh(dm_n + r * dm_nh) + h = (1.0 - z) * n + z * h + + h_offsets = (batch * seq_len + time) * hidden_size + offs_h + tl.store(out_ptr + h_offsets, h, mask=mask_h) + if save_intermediates: + tl.store(delta_x_ptr + x_offsets, dx, mask=mask_i) + tl.store(delta_h_ptr + h_offsets, dh, mask=mask_h) + tl.store(gate_r_ptr + h_offsets, r, mask=mask_h) + tl.store(gate_z_ptr + h_offsets, z, mask=mask_h) + tl.store(gate_n_ptr + h_offsets, n, mask=mask_h) + tl.store(dm_nh_hist_ptr + h_offsets, dm_nh, mask=mask_h) + tl.store(mask_x_ptr + x_offsets, update_x, mask=mask_i) + tl.store(mask_h_ptr + h_offsets, update_h, mask=mask_h) + + if collect_stats: + zeros_x += tl.sum(((dx == 0.0) & mask_i).to(tl.int64), axis=0) + zeros_h += tl.sum(((dh == 0.0) & mask_h).to(tl.int64), axis=0) + + if collect_stats: + tl.atomic_add(stats_ptr, zeros_x) + tl.atomic_add(stats_ptr + 1, zeros_h) + + @triton.jit + def _deltagru_backward_kernel( + grad_out_ptr, + h0_ptr, + out_ptr, + delta_x_ptr, + delta_h_ptr, + gate_r_ptr, + gate_z_ptr, + gate_n_ptr, + dm_nh_hist_ptr, + mask_x_ptr, + mask_h_ptr, + weight_x_ptr, + weight_h_ptr, + grad_x_ptr, + grad_x_p0_ptr, + grad_h0_ptr, + grad_h_p0_ptr, + grad_dm_nh0_ptr, + grad_dm0_ptr, + grad_mac_x_ptr, + grad_mac_h_ptr, + seq_len, + input_size: tl.constexpr, + hidden_size: tl.constexpr, + block_i: tl.constexpr, + block_h: tl.constexpr, + ): + batch = tl.program_id(0) + offs_i = tl.arange(0, block_i) + offs_h = tl.arange(0, block_h) + mask_i = offs_i < input_size + mask_h = offs_h < hidden_size + + wx_offsets = offs_h[:, None] * input_size + offs_i[None, :] + wh_offsets = offs_h[:, None] * hidden_size + offs_h[None, :] + wx_r = tl.load(weight_x_ptr + wx_offsets, mask=mask_h[:, None] & mask_i[None, :], other=0.0) + wx_z = tl.load( + weight_x_ptr + hidden_size * input_size + wx_offsets, + mask=mask_h[:, None] & mask_i[None, :], + other=0.0, + ) + wx_n = tl.load( + weight_x_ptr + 2 * hidden_size * input_size + wx_offsets, + mask=mask_h[:, None] & mask_i[None, :], + other=0.0, + ) + wh_r = tl.load(weight_h_ptr + wh_offsets, mask=mask_h[:, None] & mask_h[None, :], other=0.0) + wh_z = tl.load( + weight_h_ptr + hidden_size * hidden_size + wh_offsets, + mask=mask_h[:, None] & mask_h[None, :], + other=0.0, + ) + wh_n = tl.load( + weight_h_ptr + 2 * hidden_size * hidden_size + wh_offsets, + mask=mask_h[:, None] & mask_h[None, :], + other=0.0, + ) + + grad_x_p = tl.zeros((block_i,), tl.float32) + grad_h = tl.zeros((block_h,), tl.float32) + grad_h_p = tl.zeros((block_h,), tl.float32) + grad_dm_r = tl.zeros((block_h,), tl.float32) + grad_dm_z = tl.zeros((block_h,), tl.float32) + grad_dm_n = tl.zeros((block_h,), tl.float32) + grad_dm_nh = tl.zeros((block_h,), tl.float32) + h0 = tl.load(h0_ptr + batch * hidden_size + offs_h, mask=mask_h, other=0.0) + + for reverse_time in tl.range(0, seq_len, loop_unroll_factor=1): + time = seq_len - 1 - reverse_time + h_offsets = (batch * seq_len + time) * hidden_size + offs_h + x_offsets = (batch * seq_len + time) * input_size + offs_i + + grad_h += tl.load(grad_out_ptr + h_offsets, mask=mask_h, other=0.0) + h_prev = tl.load( + out_ptr + h_offsets - hidden_size, + mask=mask_h & (time > 0), + other=0.0, + ) + h_prev = tl.where(time > 0, h_prev, h0) + r = tl.load(gate_r_ptr + h_offsets, mask=mask_h, other=0.0) + z = tl.load(gate_z_ptr + h_offsets, mask=mask_h, other=0.0) + n = tl.load(gate_n_ptr + h_offsets, mask=mask_h, other=0.0) + dm_nh = tl.load(dm_nh_hist_ptr + h_offsets, mask=mask_h, other=0.0) + + grad_n = grad_h * (1.0 - z) + grad_z = grad_h * (h_prev - n) + grad_h_prev = grad_h * z + grad_pre_n = grad_n * (1.0 - n * n) + grad_r = grad_pre_n * dm_nh + grad_dm_nh += grad_pre_n * r + grad_dm_r += grad_r * r * (1.0 - r) + grad_dm_z += grad_z * z * (1.0 - z) + grad_dm_n += grad_pre_n + + grad_mac_x_r = grad_dm_r + grad_mac_x_z = grad_dm_z + grad_mac_x_n = grad_dm_n + grad_mac_h_r = grad_dm_r + grad_mac_h_z = grad_dm_z + grad_mac_h_n = grad_dm_nh + + grad_dx = ( + tl.sum(wx_r * grad_mac_x_r[:, None], axis=0) + + tl.sum(wx_z * grad_mac_x_z[:, None], axis=0) + + tl.sum(wx_n * grad_mac_x_n[:, None], axis=0) + ) + grad_dh = ( + tl.sum(wh_r * grad_mac_h_r[:, None], axis=0) + + tl.sum(wh_z * grad_mac_h_z[:, None], axis=0) + + tl.sum(wh_n * grad_mac_h_n[:, None], axis=0) + ) + + update_x = tl.load(mask_x_ptr + x_offsets, mask=mask_i, other=0).to(tl.int1) + update_h = tl.load(mask_h_ptr + h_offsets, mask=mask_h, other=0).to(tl.int1) + saved_delta_x = tl.load(delta_x_ptr + x_offsets, mask=mask_i, other=0.0) + saved_delta_h = tl.load(delta_h_ptr + h_offsets, mask=mask_h, other=0.0) + # For non-NaN thresholds, masked-fill's keep condition differs from + # the state-update condition only when a raw delta is NaN. + keep_delta_x = update_x | (saved_delta_x != saved_delta_x) + keep_delta_h = update_h | (saved_delta_h != saved_delta_h) + grad_xt = tl.where(keep_delta_x, grad_dx, 0.0) + tl.where( + update_x, grad_x_p, 0.0 + ) + grad_x_p = tl.where(keep_delta_x, -grad_dx, 0.0) + tl.where( + update_x, 0.0, grad_x_p + ) + grad_h_prev += tl.where(keep_delta_h, grad_dh, 0.0) + tl.where( + update_h, grad_h_p, 0.0 + ) + grad_h_p = tl.where(keep_delta_h, -grad_dh, 0.0) + tl.where( + update_h, 0.0, grad_h_p + ) + + tl.store(grad_x_ptr + x_offsets, grad_xt, mask=mask_i) + mac_offsets = (batch * seq_len + time) * (3 * hidden_size) + offs_h + tl.store(grad_mac_x_ptr + mac_offsets, grad_mac_x_r, mask=mask_h) + tl.store( + grad_mac_x_ptr + mac_offsets + hidden_size, + grad_mac_x_z, + mask=mask_h, + ) + tl.store( + grad_mac_x_ptr + mac_offsets + 2 * hidden_size, + grad_mac_x_n, + mask=mask_h, + ) + tl.store(grad_mac_h_ptr + mac_offsets, grad_mac_h_r, mask=mask_h) + tl.store( + grad_mac_h_ptr + mac_offsets + hidden_size, + grad_mac_h_z, + mask=mask_h, + ) + tl.store( + grad_mac_h_ptr + mac_offsets + 2 * hidden_size, + grad_mac_h_n, + mask=mask_h, + ) + + # Gradients for the accumulator carries at the previous timestep. + grad_dm_r = grad_mac_x_r + grad_dm_z = grad_mac_x_z + grad_dm_n = grad_mac_x_n + grad_dm_nh = grad_mac_h_n + grad_h = grad_h_prev + + tl.store(grad_x_p0_ptr + batch * input_size + offs_i, grad_x_p, mask=mask_i) + tl.store(grad_h0_ptr + batch * hidden_size + offs_h, grad_h, mask=mask_h) + tl.store(grad_h_p0_ptr + batch * hidden_size + offs_h, grad_h_p, mask=mask_h) + tl.store(grad_dm_nh0_ptr + batch * hidden_size + offs_h, grad_dm_nh, mask=mask_h) + dm0_offsets = batch * (3 * hidden_size) + offs_h + tl.store(grad_dm0_ptr + dm0_offsets, grad_dm_r, mask=mask_h) + tl.store(grad_dm0_ptr + dm0_offsets + hidden_size, grad_dm_z, mask=mask_h) + tl.store(grad_dm0_ptr + dm0_offsets + 2 * hidden_size, grad_dm_n, mask=mask_h) + + +class _TritonDeltaGRUFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + x_p0: torch.Tensor, + h0: torch.Tensor, + h_p0: torch.Tensor, + dm_nh0: torch.Tensor, + dm0: torch.Tensor, + weight_x: torch.Tensor, + weight_h: torch.Tensor, + stats: torch.Tensor, + th_x: float, + th_h: float, + collect_stats: bool, + ) -> torch.Tensor: + batch_size, seq_len, input_size = input.shape + hidden_size = h0.shape[-1] + output = input.new_empty((batch_size, seq_len, hidden_size)) + delta_x = torch.empty_like(input) + state_shape = (batch_size, seq_len, hidden_size) + delta_h = input.new_empty(state_shape) + gate_r = input.new_empty(state_shape) + gate_z = input.new_empty(state_shape) + gate_n = input.new_empty(state_shape) + dm_nh_hist = input.new_empty(state_shape) + mask_x = torch.empty_like(input, dtype=torch.uint8) + mask_h = torch.empty(state_shape, device=input.device, dtype=torch.uint8) + + block_i = triton.next_power_of_2(input_size) + block_h = triton.next_power_of_2(hidden_size) + _deltagru_forward_kernel[(batch_size,)]( + input, + x_p0, + h0, + h_p0, + dm_nh0, + dm0, + weight_x, + weight_h, + output, + delta_x, + delta_h, + gate_r, + gate_z, + gate_n, + dm_nh_hist, + mask_x, + mask_h, + stats, + seq_len=seq_len, + input_size=input_size, + hidden_size=hidden_size, + th_x=th_x, + th_h=th_h, + collect_stats=collect_stats, + save_intermediates=True, + block_i=block_i, + block_h=block_h, + num_warps=1, + ) + ctx.save_for_backward( + h0, + output, + delta_x, + delta_h, + gate_r, + gate_z, + gate_n, + dm_nh_hist, + mask_x, + mask_h, + weight_x, + weight_h, + ) + ctx.input_size = input_size + ctx.hidden_size = hidden_size + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if torch.is_grad_enabled(): + raise RuntimeError( + "fused DeltaGRU supports first-order training only; set " + "OPENDPD_DISABLE_TRITON_DELTAGRU=1 for higher-order autograd" + ) + ( + h0, + output, + delta_x, + delta_h, + gate_r, + gate_z, + gate_n, + dm_nh_hist, + mask_x, + mask_h, + weight_x, + weight_h, + ) = ctx.saved_tensors + grad_output = grad_output.contiguous() + batch_size, seq_len, hidden_size = output.shape + input_size = ctx.input_size + + grad_input = delta_x.new_empty(delta_x.shape) + grad_x_p0 = delta_x.new_empty((batch_size, input_size)) + grad_h0 = output.new_empty((batch_size, hidden_size)) + grad_h_p0 = output.new_empty((batch_size, hidden_size)) + grad_dm_nh0 = output.new_empty((batch_size, hidden_size)) + grad_dm0 = output.new_empty((batch_size, 3 * hidden_size)) + grad_mac_x = output.new_empty((batch_size, seq_len, 3 * hidden_size)) + grad_mac_h = output.new_empty((batch_size, seq_len, 3 * hidden_size)) + + block_i = triton.next_power_of_2(input_size) + block_h = triton.next_power_of_2(hidden_size) + _deltagru_backward_kernel[(batch_size,)]( + grad_output, + h0, + output, + delta_x, + delta_h, + gate_r, + gate_z, + gate_n, + dm_nh_hist, + mask_x, + mask_h, + weight_x, + weight_h, + grad_input, + grad_x_p0, + grad_h0, + grad_h_p0, + grad_dm_nh0, + grad_dm0, + grad_mac_x, + grad_mac_h, + seq_len=seq_len, + input_size=input_size, + hidden_size=hidden_size, + block_i=block_i, + block_h=block_h, + num_warps=1, + ) + + flat_delta_x = delta_x.reshape(-1, input_size) + flat_delta_h = delta_h.reshape(-1, hidden_size) + flat_grad_mac_x = grad_mac_x.reshape(-1, 3 * hidden_size) + flat_grad_mac_h = grad_mac_h.reshape(-1, 3 * hidden_size) + grad_weight_x = flat_grad_mac_x.t().matmul(flat_delta_x) + grad_weight_h = flat_grad_mac_h.t().matmul(flat_delta_h) + + return ( + grad_input, + grad_x_p0, + grad_h0, + grad_h_p0, + grad_dm_nh0, + grad_dm0, + grad_weight_x, + grad_weight_h, + None, + None, + None, + None, + ) + + +def triton_deltagru( + input: torch.Tensor, + x_p0: torch.Tensor, + h0: torch.Tensor, + h_p0: torch.Tensor, + dm_nh0: torch.Tensor, + dm0: torch.Tensor, + weight_x: torch.Tensor, + weight_h: torch.Tensor, + stats: torch.Tensor, + th_x: float, + th_h: float, + collect_stats: bool = False, +) -> torch.Tensor: + """Run the fused recurrence on a contiguous ``(batch, time, input)`` tensor.""" + + if triton is None: # pragma: no cover + raise RuntimeError("Triton is not installed") + if math.isnan(th_x) or math.isnan(th_h): + raise ValueError("DeltaGRU thresholds must not be NaN") + batch_size, seq_len, input_size = input.shape + hidden_size = h0.shape[-1] if h0.ndim else 0 + expected = ( + (x_p0, (batch_size, input_size)), + (h0, (batch_size, hidden_size)), + (h_p0, (batch_size, hidden_size)), + (dm_nh0, (batch_size, hidden_size)), + (dm0, (batch_size, 3 * hidden_size)), + (weight_x, (3 * hidden_size, input_size)), + (weight_h, (3 * hidden_size, hidden_size)), + ) + if not can_use_triton_deltagru(input, hidden_size): + raise ValueError("unsupported DeltaGRU input tensor or hidden size") + if not all( + tensor.shape == shape + and tensor.device == input.device + and tensor.dtype == input.dtype + and tensor.is_contiguous() + for tensor, shape in expected + ): + raise ValueError("DeltaGRU states and weights must match input shape/device/dtype") + if ( + stats.device != input.device + or stats.dtype != torch.int64 + or stats.numel() < 2 + or not stats.is_contiguous() + ): + raise ValueError("stats must be a contiguous CUDA int64 tensor with at least 2 elements") + if not torch.is_grad_enabled(): + output = input.new_empty((batch_size, seq_len, hidden_size)) + dummy = input.new_empty(1) + dummy_mask = torch.empty(1, device=input.device, dtype=torch.uint8) + _deltagru_forward_kernel[(batch_size,)]( + input, + x_p0, + h0, + h_p0, + dm_nh0, + dm0, + weight_x, + weight_h, + output, + dummy, + dummy, + dummy, + dummy, + dummy, + dummy, + dummy_mask, + dummy_mask, + stats, + seq_len=seq_len, + input_size=input_size, + hidden_size=hidden_size, + th_x=th_x, + th_h=th_h, + collect_stats=collect_stats, + save_intermediates=False, + block_i=triton.next_power_of_2(input_size), + block_h=triton.next_power_of_2(hidden_size), + num_warps=1, + ) + return output + return _TritonDeltaGRUFunction.apply( + input, + x_p0, + h0, + h_p0, + dm_nh0, + dm0, + weight_x, + weight_h, + stats, + th_x, + th_h, + collect_stats, + ) diff --git a/bash_scripts/OpenDPDv2.sh b/bash_scripts/OpenDPDv2.sh index f5cf179..27c4aa5 100644 --- a/bash_scripts/OpenDPDv2.sh +++ b/bash_scripts/OpenDPDv2.sh @@ -40,6 +40,11 @@ lr=${LR:-5e-3} lr_end=${LR_END:-1e-4} decay_factor=${DECAY_FACTOR:-0.5} patience=${PATIENCE:-10} +cuda_graph_training=${CUDA_GRAPH_TRAINING:-1} +cuda_graph_opts=() +if [[ "${cuda_graph_training}" == "1" ]]; then + cuda_graph_opts+=(--cuda_graph_training) +fi seed_values=(0) PA_backbone=dgru @@ -107,7 +112,8 @@ for i_seed in "${seed_values[@]}"; do --decay_factor "${decay_factor}" \ --patience "${patience}" \ --thx "${thx}" \ - --thh "${thh}" + --thh "${thh}" \ + "${cuda_graph_opts[@]}" quant_dir_label="w${quant_n_bits_w}a${quant_n_bits_a}" diff --git a/benchmark/benchmark_delta_dpd_results.png b/benchmark/benchmark_delta_dpd_results.png new file mode 100644 index 0000000..80647f6 Binary files /dev/null and b/benchmark/benchmark_delta_dpd_results.png differ diff --git a/benchmark/benchmark_report.md b/benchmark/benchmark_report.md index 923d10f..efb3239 100644 --- a/benchmark/benchmark_report.md +++ b/benchmark/benchmark_report.md @@ -1,187 +1,138 @@ -# DPD Benchmark Report: AI vs. Traditional Approaches +# OpenDPD PA Modeling and DPD Benchmark -## 1. Overview - -This report benchmarks Digital Pre-Distortion (DPD) algorithms across two categories: - -- **AI-based (neural network)**: GRU and TRes-DeltaGRU, trained via stochastic gradient descent (SGD) -- **Traditional (polynomial)**: Memory Polynomial (MP) and Generalized Memory Polynomial (GMP), identified via QR decomposition (closed-form least squares) - -All models are constrained to approximately 500 real-valued parameters for a fair comparison. Evaluation uses Adjacent Channel Leakage Ratio (ACLR) and Error Vector Magnitude (EVM), where more negative values indicate better performance. +## Technical summary -## 2. Test Signals +This benchmark evaluates MP, GMP, GRU, TRes-GRU, and TRes-DeltaGRU (THX=THH=0) for PA modeling on APA_200MHz and DPA_160MHz. The DPD comparison evaluates MP, GMP, GRU, and TRes-GRU. PA models use approximately 2,700 real parameters; DPD models use approximately 1,000. Every neural run uses the same 300-epoch optimization recipe. MP PA uses direct least squares and GMP PA uses rank-controlled truncated SVD; their predistorters use the indirect learning architecture (ILA). The four-model DPD comparison uses the dataset-specific TRes-GRU-H27 PA checkpoint selected by validation NMSE. A separate sensitivity experiment trains TRes-DeltaGRU-H15 DPD independently through both the TRes-GRU-H27 and zero-threshold TRes-DeltaGRU-H27 PA surrogates. -| Property | APA_200MHz | DPA_160MHz | -|----------|-----------|-----------| -| Standard | LTE TM3.1a | OFDM | -| Configuration | 5-carrier x 40 MHz | 4-carrier x 40 MHz | -| Total Bandwidth | 200 MHz | 160 MHz | -| Modulation | 256-QAM | 1024-QAM | -| PAPR | 10.01 dB (at CCDF of 0.001%) | 10.38 dB | -| Sampling Rate | 983.04 MHz | 640 MHz | -| Dataset Size | 98,304 samples | 491,520 samples | -| Segment Size (nperseg) | 19,662 | 16,384 | -| Sub-channels | 5 | 4 | -| Dataset Split | 60% / 20% / 20% (train / val / test) | 60% / 20% / 20% (train / val / test) | - -## 3. Power Amplifier Devices Under Test (DUT) - -| Property | APA_200MHz | DPA_160MHz | -|----------|-----------|-----------| -| PA Type | GaN Doherty PA | 40 nm CMOS Digital PA (DPA) | -| Part / Technology | Ampleon AR211132 (evaluation board) | 40 nm CMOS | -| Carrier Frequency | 3.5 GHz | 2.4 GHz | -| Average Output Power | 41.2 dBm | 13.75 dBm | -| P1dB Compression Point | 46.5 dBm | - | -| P3dB Compression Point | 50 dBm | - | -| Flat Gain Bandwidth | 200 MHz | 160 MHz | -| Reference | Wu et al., "OpenDPDv2," arXiv:2507.06849 | Wu et al., "MP-DPD," IEEE IMS 2024, arXiv:2404.15364 | +## Key findings -The APA_200MHz dataset uses a high-power 3.5 GHz GaN Doherty PA designed for 5G base station applications. The DPA_160MHz dataset uses a low-power 2.4 GHz CMOS digital PA targeting Wi-Fi/IoT transmitter applications. Both PAs exhibit nonlinear distortion with memory effects, requiring DPD to meet spectral emission standards. +- **APA_200MHz:** the validation PA leader is TRes-DeltaGRU-H27 (THX=THH=0) at -39.0929 dB NMSE (-39.1777 dB test). The validation DPD leader is TRes-GRU-H15 at -53.4736 dB ACLR average (-53.4879 dB test). For TRes-DeltaGRU-H15 DPD, the better validation result uses the TRes-DeltaGRU-H27 (THX=THH=0) PA surrogate at -54.4874 dB ACLR average (-54.4010 dB test). +- **DPA_160MHz:** the validation PA leader is TRes-DeltaGRU-H27 (THX=THH=0) at -39.4802 dB NMSE (-39.7353 dB test). The validation DPD leader is TRes-GRU-H15 at -60.7425 dB ACLR average (-60.3717 dB test). For TRes-DeltaGRU-H15 DPD, the better validation result uses the TRes-DeltaGRU-H27 (THX=THH=0) PA surrogate at -61.9253 dB ACLR average (-61.6586 dB test). +- **APA GMP stability:** column-normalized truncated SVD at `rcond=1e-04` retains 650/1,350 singular directions. Validation/test NMSE is -38.70/-38.66 dB. The fixed cutoff suppresses ill-conditioned delayed-envelope directions instead of applying CUDA `gels`'s invalid full-rank assumption. -## 4. PA Surrogate Model +![Test-set PA modeling and DPD results](benchmark_results.png) -All DPD models are evaluated through the same trained GRU PA surrogate model per dataset, ensuring a fair comparison. The PA model is frozen during DPD training (indirect learning architecture). +*Test split; more negative is better. Outlined points are the models selected by validation. DPD is simulated through the dataset-specific TRes-GRU-H27 PA surrogate.* -| Dataset | PA Backbone | PA Hidden Size | PA Parameters | PA Training | Best PA NMSE | -|---------|-------------|----------------|---------------|-------------|-------------| -| APA_200MHz | GRU | 23 | 1,911 | 100 epochs, AdamW, lr=5e-4 | -43.52 dB | -| DPA_160MHz | GRU | 24 | 2,066 | 100 epochs, AdamW, lr=5e-4 | -38.43 dB | +![TRes-DeltaGRU DPD PA-surrogate sensitivity](benchmark_delta_dpd_results.png) -## 5. DPD Algorithms +*Test split; each point is an independently trained TRes-DeltaGRU-H15 DPD using the named frozen PA surrogate. These results compare surrogate sensitivity, not measurements from one shared physical PA.* -### 5.1 AI-Based Models (trained via SGD) +## APA_200MHz -**GRU (Gated Recurrent Unit)** -- Standard PyTorch GRU cell followed by a linear output layer -- Input: I/Q samples (2D); Output: I/Q samples (2D) -- Configuration: hidden_size=11, num_layers=1 -- Parameters: 519 real-valued +### PA modeling -**TRes-DeltaGRU (Temporal Residual Delta GRU)** -- Delta GRU cell with TCN-based temporal residual skip connection (Wu et al., OpenDPDv2, arXiv:2507.06849) -- The TCN residual path decouples output dynamics from hidden-state sparsity, enabling aggressive temporal sparsity without linearization loss -- Configuration: hidden_size=10, num_layers=1, Conv1d(2,3,k=3,d=16) + Conv1d(3,2,k=1) -- Parameters: 524 real-valued +| Model | Parameters | Method | Selected epoch | NMSE, val / test (dB) | EVM, val / test (dB) | Validation ACLR L / R / avg (dB) | Test ACLR L / R / avg (dB) | +|---|---:|---|---:|---:|---:|---:|---:| +| MP | 2,700 | Direct least squares | N/A | -37.0405 / -36.9635 | -40.2302 / -39.9571 | -27.4060 / -27.3703 / -27.3881 | -27.8869 / -27.7765 / -27.8317 | +| GMP | 2,700 | Truncated SVD (rank 650/1,350) | N/A | -38.7020 / -38.6606 | -43.0540 / -42.7971 | -27.4053 / -27.4129 / -27.4091 | -27.7830 / -27.6294 / -27.7062 | +| GRU-H28 | 2,746 | Supervised, AdamW | 292 | -38.8504 / -38.9367 | -43.6317 / -43.5732 | -27.3455 / -27.4211 / -27.3833 | -27.6783 / -27.6725 / -27.6754 | +| TRes-GRU-H27 | 2,751 | Supervised, AdamW | 288 | -39.0447 / -39.1293 | -43.9118 / -43.9299 | -27.4100 / -27.4074 / -27.4087 | -27.7503 / -27.6404 / -27.6953 | +| TRes-DeltaGRU-H27 (THX=THH=0) | 2,751 | Supervised, AdamW | 294 | -39.0929 / -39.1777 | -44.0254 / -43.9683 | -27.3872 / -27.4456 / -27.4164 | -27.6800 / -27.6721 / -27.6761 | -### 5.2 Traditional Models (identified via QR) +### DPD -**MP (Memory Polynomial)** -- Diagonal subset of the Volterra series: basis functions `x(n-q) * |x(n-q)|^k` -- Configuration: K=5 nonlinearity orders, Q=50 memory depth -- Parameters: 250 complex coefficients = 500 real-valued -- Identification: one-shot least-squares via `numpy.linalg.lstsq` (SVD-based) +| Model | Parameters | Method | Selected epoch | NMSE, val / test (dB) | EVM, val / test (dB) | Validation ACLR L / R / avg (dB) | Test ACLR L / R / avg (dB) | +|---|---:|---|---:|---:|---:|---:|---:| +| MP | 1,000 | ILA, least squares | N/A | -42.1149 / -42.1896 | -48.1895 / -48.1535 | -45.5590 / -44.2857 / -44.9224 | -46.0740 / -44.2980 / -45.1860 | +| GMP | 1,000 | ILA, least squares | N/A | -38.3167 / -38.5259 | -46.2753 / -46.3523 | -43.5560 / -42.3895 / -42.9727 | -44.0657 / -43.1207 / -43.5932 | +| GRU-H16 | 994 | DLA, AdamW | 299 | -45.1735 / -45.1315 | -47.6885 / -47.4275 | -51.3857 / -51.2623 / -51.3240 | -51.0559 / -50.9669 / -51.0114 | +| TRes-GRU-H15 | 999 | DLA, AdamW | 296 | -44.0368 / -44.2859 | -44.8654 / -45.0973 | -53.2065 / -53.7407 / -53.4736 | -53.4693 / -53.5066 / -53.4879 | -**GMP (Generalized Memory Polynomial)** -- Extended Memory Polynomial with lagging and leading cross-terms (Morgan et al., IEEE TSP, 2006) -- Aligned terms: `x(n-q) * |x(n-q)|^k` (Ka=5, La=15 -> 75 terms) -- Lagging cross-terms: `x(n-q) * |x(n-q-l)|^k` (Kb=4, Lb=15, Mb=2 -> 120 terms) -- Leading cross-terms: `x(n-q) * |x(n-q+l)|^k` (Kc=4, Lc=15, Mc=1 -> 60 terms) -- Parameters: 255 complex coefficients = 510 real-valued -- Identification: one-shot least-squares via `numpy.linalg.lstsq` (SVD-based) +### TRes-DeltaGRU DPD by PA surrogate -Both traditional models use the Indirect Learning Architecture (ILA): build basis from normalized PA output, solve for postdistorter coefficients, then copy to predistorter. +Both rows use TRes-DeltaGRU-H15 with 999 parameters and THX=THH=0. They are separately trained through the indicated frozen PA surrogate. -## 6. Training Configuration +| PA surrogate | PA NMSE, val / test (dB) | DPD parameters | Selected epoch | DPD NMSE, val / test (dB) | DPD EVM, val / test (dB) | DPD ACLR avg, val / test (dB) | +|---|---:|---:|---:|---:|---:|---:| +| TRes-GRU-H27 | -39.0447 / -39.1293 | 999 | 264 | -46.8471 / -46.9184 | -48.3486 / -48.2629 | -53.2898 / -53.6435 | +| TRes-DeltaGRU-H27 (THX=THH=0) | -39.0929 / -39.1777 | 999 | 296 | -47.8582 / -47.9703 | -49.8200 / -49.6999 | -54.4874 / -54.4010 | -### SGD-Based Models (GRU, TRes-DeltaGRU) +## DPA_160MHz -| Parameter | Value | -|-----------|-------| -| Optimizer | AdamW | -| Learning Rate | 5e-4 | -| Epochs | 300 | -| Batch Size | 256 | -| Loss Function | MSE (L2) | -| LR Schedule | None | -| Gradient Clipping | 200 | -| Best Model Selection | Validation ACLR_AVG | -| Seed | 0 | +### PA modeling -### QR-Based Models (MP, GMP) +| Model | Parameters | Method | Selected epoch | NMSE, val / test (dB) | EVM, val / test (dB) | Validation ACLR L / R / avg (dB) | Test ACLR L / R / avg (dB) | +|---|---:|---|---:|---:|---:|---:|---:| +| MP | 2,700 | Direct least squares | N/A | -37.9823 / -38.3506 | -40.3940 / -40.9035 | -34.9319 / -34.4718 / -34.7018 | -35.1443 / -34.6409 / -34.8926 | +| GMP | 2,700 | Truncated SVD (rank 918/1,350) | N/A | -38.9088 / -39.4014 | -41.6537 / -42.3839 | -34.9108 / -34.3514 / -34.6311 | -35.0859 / -34.5171 / -34.8015 | +| GRU-H28 | 2,746 | Supervised, AdamW | 34 | -39.0801 / -39.4908 | -42.1993 / -42.8282 | -34.9599 / -34.5649 / -34.7624 | -35.0367 / -34.6393 / -34.8380 | +| TRes-GRU-H27 | 2,751 | Supervised, AdamW | 25 | -39.3229 / -39.7044 | -42.5812 / -43.1447 | -34.9905 / -34.4694 / -34.7299 | -35.1550 / -34.6605 / -34.9078 | +| TRes-DeltaGRU-H27 (THX=THH=0) | 2,751 | Supervised, AdamW | 16 | -39.4802 / -39.7353 | -42.9642 / -43.3196 | -34.9649 / -34.4466 / -34.7057 | -35.1432 / -34.6450 / -34.8941 | -| Parameter | Value | -|-----------|-------| -| Solver | `numpy.linalg.lstsq` (SVD/QR internally) | -| Training Iterations | 1 (closed-form) | -| Regularization | None | -| Training Data | Full training set | +### DPD -## 7. Results +| Model | Parameters | Method | Selected epoch | NMSE, val / test (dB) | EVM, val / test (dB) | Validation ACLR L / R / avg (dB) | Test ACLR L / R / avg (dB) | +|---|---:|---|---:|---:|---:|---:|---:| +| MP | 1,000 | ILA, least squares | N/A | -43.5581 / -43.1938 | -50.4204 / -50.5002 | -49.0670 / -52.1969 / -50.6320 | -49.4442 / -52.4261 / -50.9351 | +| GMP | 1,000 | ILA, least squares | N/A | -44.2662 / -43.9721 | -50.1286 / -50.1887 | -52.3572 / -53.2123 / -52.7847 | -52.6167 / -53.6083 / -53.1125 | +| GRU-H16 | 994 | DLA, AdamW | 297 | -49.9059 / -49.7602 | -53.7314 / -53.9790 | -57.6084 / -57.0975 / -57.3530 | -57.6760 / -56.8525 / -57.2642 | +| TRes-GRU-H15 | 999 | DLA, AdamW | 296 | -52.9106 / -53.1938 | -57.1425 / -57.8451 | -60.4579 / -61.0271 / -60.7425 | -60.2712 / -60.4722 / -60.3717 | -### 7.1 APA_200MHz (3.5 GHz GaN Doherty PA, 5-carrier LTE, 200 MHz BW, 256-QAM) +### TRes-DeltaGRU DPD by PA surrogate -| Rank | Model | Parameters | Method | ACLR_AVG (dB) | EVM (dB) | -|------|-------|-----------|--------|---------------|----------| -| 1 | **TRes-DeltaGRU** | 524 | SGD 300ep | **-53.35** | **-49.08** | -| 2 | GRU | 519 | SGD 300ep | -52.61 | -47.50 | -| 3 | MP | 500 | QR | -41.01 | -32.68 | -| 4 | GMP | 510 | QR | -38.80 | -38.53 | +Both rows use TRes-DeltaGRU-H15 with 999 parameters and THX=THH=0. They are separately trained through the indicated frozen PA surrogate. -### 7.2 DPA_160MHz (2.4 GHz CMOS Digital PA, 4-carrier, 160 MHz BW, 1024-QAM) +| PA surrogate | PA NMSE, val / test (dB) | DPD parameters | Selected epoch | DPD NMSE, val / test (dB) | DPD EVM, val / test (dB) | DPD ACLR avg, val / test (dB) | +|---|---:|---:|---:|---:|---:|---:| +| TRes-GRU-H27 | -39.3229 / -39.7044 | 999 | 295 | -53.8684 / -54.2446 | -57.9312 / -58.8486 | -61.5425 / -61.0442 | +| TRes-DeltaGRU-H27 (THX=THH=0) | -39.4802 / -39.7353 | 999 | 293 | -54.1159 / -54.3722 | -58.2478 / -59.3978 | -61.9253 / -61.6586 | -| Rank | Model | Parameters | Method | ACLR_AVG (dB) | EVM (dB) | -|------|-------|-----------|--------|---------------|----------| -| 1 | **TRes-DeltaGRU** | 524 | SGD 300ep | **-56.81** | **-54.00** | -| 2 | GMP | 510 | QR | -54.02 | -51.08 | -| 3 | GRU | 519 | SGD 300ep | -51.93 | -49.97 | -| 4 | MP | 500 | QR | -51.29 | -50.26 | +## Definitions -## 8. Analysis +- **NMSE:** normalized mean-square error in dB; more negative is better. The implementation averages per-segment dB ratios rather than pooling all samples into one ratio. +- **EVM:** the repository-specific mean absolute complex-spectrum error within the configured main channel, normalized within each subchannel by reference-spectrum magnitude and converted with `20 log10`. It is not demodulated constellation EVM; more negative is better. +- **ACLR L / R / avg:** adjacent-subchannel power normalized by the strongest configured main-channel subchannel, plus the arithmetic mean of left and right, in dB. More negative is better. +- **Parameters:** neural checkpoint tensor elements. MP/GMP count two real degrees of freedom for every complex coefficient. +- **Selected epoch:** zero-based neural checkpoint epoch. It is N/A for closed-form least-squares fits. -### AI vs. Traditional +## Model configurations -- On **APA_200MHz** (wider bandwidth, more carriers), the AI models dominate: TRes-DeltaGRU leads the best traditional method (MP) by **12.3 dB in ACLR** and **10.6 dB in EVM**. The recurrent hidden state captures complex long-range memory effects that polynomial basis functions cannot. -- On **DPA_160MHz**, the gap narrows but AI still leads: TRes-DeltaGRU beats the best traditional method (GMP/QR) by **2.8 dB in ACLR** and **2.9 dB in EVM**. The traditional GMP/QR is competitive here, outperforming vanilla GRU. +| Stage | MP | GMP | GRU | TRes-GRU | TRes-DeltaGRU | +|---|---|---|---|---|---| +| PA modeling | K=9, Q=150 | Ka/La=5/30; Kb/Lb/Mb=4/30/5; Kc/Lc/Mc=4/30/5 | H28 | H27 | H27, THX=THH=0 | +| DPD | K=5, Q=100 | Ka/La=5/20; Kb/Lb/Mb=4/20/3; Kc/Lc/Mc=4/20/2 | H16 | H15 | H15, THX=THH=0 (PA-surrogate sensitivity) | -### QR vs. SGD for Polynomial Models +## Temporal context and sequence boundaries -The identification method matters enormously for polynomial models. Comparing GMP trained with QR vs. SGD (100 epochs, from earlier experiments): +In sequence interiors, PA GMP uses up to five future samples and DPD GMP uses up to two through their leading-envelope terms. TRes-GRU and TRes-DeltaGRU use one-sample right context in their recurrent features and 16-sample right context in their dilated residual convolution. MP and GRU use no explicit future samples. -| Dataset | GMP (QR) ACLR | GMP (SGD 100ep) ACLR | Improvement | -|---------|--------------|---------------------|------------| -| APA_200MHz | -38.80 dB | -29.67 dB | +9.1 dB | -| DPA_160MHz | -54.02 dB | -44.46 dB | +9.6 dB | +These are offline segmented evaluations. GMP delay accesses are zero-filled and reset at each `nperseg` boundary. In both TRes models, `torch.roll(..., shifts=-1)` wraps the final position to the first sample of the same supplied sequence; the convolution zero-pads both boundaries, and recurrent state resets for each sequence. Neural optimization uses overlapping 200-sample frames with stride 1, while validation and test use independent `nperseg` segments. -QR finds the global optimum in one shot; SGD with 100 epochs of Adam cannot reach it for this convex problem. This demonstrates that when the model is linear-in-parameters, closed-form solvers are strictly superior to iterative gradient descent. +## Methodology -### TRes-DeltaGRU vs. GRU +Neural PA and DPD models use batch size 64, 300 epochs, AdamW with MSE loss, initial learning rate 5e-3, and ReduceLROnPlateau with factor 0.5, patience 5, and minimum learning rate 5e-5. Frame length is 200, frame stride is 1, and seed is 0. PA checkpoints are selected by minimum validation NMSE; neural DPD checkpoints are selected by minimum validation ACLR average. -TRes-DeltaGRU consistently outperforms vanilla GRU across both datasets: +The TRes-DeltaGRU PA and DPD runs use input-delta threshold THX=0 and hidden-state-delta threshold THH=0. This disables threshold-induced temporal pruning; exact arithmetic deltas may still naturally be zero. The configuration therefore evaluates the dense zero-threshold recurrence, not a sparsity or efficiency claim. -| Dataset | TRes-DeltaGRU ACLR | GRU ACLR | Gap | -|---------|-------------------|----------|-----| -| APA_200MHz | -53.35 dB | -52.61 dB | +0.7 dB | -| DPA_160MHz | -56.81 dB | -51.93 dB | +4.9 dB | +AdamW uses weight decay 0.01, betas (0.9, 0.999), and epsilon 1e-8. The scheduler uses relative threshold 1e-4, cooldown 0, and epsilon 1e-8. -The TCN skip connection provides a direct residual path from input to output, decoupling the output from hidden-state dynamics. This is especially beneficial on DPA_160MHz where the improvement is nearly 5 dB. +MP and GMP are complex polynomial models fit after L2 column scaling. MP and both ILA-DPD fits use `torch.linalg.lstsq` (`gels`). GMP PA modeling uses `torch.linalg.svd` (`gesvdj`) with a fixed relative cutoff of 1e-04; the effective retained rank is reported with each GMP PA result. These closed-form fits do not use the neural batch, epoch, optimizer, or learning-rate settings. PA polynomial fits map measured PA input to output directly. DPD polynomial fits use ILA, fitting a postdistorter and copying its coefficients to the predistorter. No ridge penalty or validation-tuned regularization is applied. -## 9. Reproducing These Results +Each dataset uses its independently trained TRes-GRU-H27 PA surrogate for the four-model DPD comparison. The two TRes-DeltaGRU-H15 sensitivity rows are separate DPD training runs, one through that TRes-GRU PA and one through the independently trained TRes-DeltaGRU-H27 PA. Training data supply all gradient and least-squares fits. Neural validation metrics drive the learning-rate schedule and checkpoint selection; test data are not used for fitting, scheduling, or selection. -### SGD Models +## Limitations and robustness -```bash -# Train PA model (once per dataset) -python main.py --step train_pa --dataset_name APA_200MHz --PA_backbone gru --PA_hidden_size 23 --n_epochs 100 -python main.py --step train_pa --dataset_name DPA_160MHz --PA_backbone gru --PA_hidden_size 24 --n_epochs 100 +- Neural execution uses soft reproducibility with cuDNN benchmark enabled, so repeated runs can differ slightly. +- DPD scores measure simulated performance through a learned PA surrogate, not a fresh over-the-air or bench measurement. +- One seed is evaluated. The reported table is not a distribution over training runs. +- The five PA candidates and four primary DPD candidates are matched approximately by real parameter count, not by FLOPs, latency, memory traffic, or fit time. The PA-surrogate sensitivity experiment adds two independently trained TRes-DeltaGRU DPD runs per dataset. +- Results obtained through different learned PA surrogates are simulator-sensitivity evidence; they are not a controlled ranking against one shared physical PA response. +- CUDA `gels` assumes a full-rank design matrix and does not return numerical rank. It remains in use for MP and ILA-DPD; the GMP PA path records its SVD spectrum, cutoff, and retained rank. +- GMP PA has 1,350 stored complex coefficients (2,700 nominal real parameters), but truncated SVD reduces its effective rank. The comparison is matched by stored coefficient count, not effective degrees of freedom. +- Look-ahead is an input dependency, not measured inference latency. Streaming reformulations and continuous boundary/state handling are not evaluated. -# Train GRU DPD -python main.py --step train_dpd --dataset_name APA_200MHz --PA_backbone gru --PA_hidden_size 23 --DPD_backbone gru --DPD_hidden_size 11 --n_epochs 300 -python main.py --step train_dpd --dataset_name DPA_160MHz --PA_backbone gru --PA_hidden_size 24 --DPD_backbone gru --DPD_hidden_size 11 --n_epochs 300 +## Provenance -# Train TRes-DeltaGRU DPD -python main.py --step train_dpd --dataset_name APA_200MHz --PA_backbone gru --PA_hidden_size 23 --DPD_backbone tres_deltagru --DPD_hidden_size 10 --n_epochs 300 -python main.py --step train_dpd --dataset_name DPA_160MHz --PA_backbone gru --PA_hidden_size 24 --DPD_backbone tres_deltagru --DPD_hidden_size 10 --n_epochs 300 -``` - -### QR Models - -```bash -# MP (K=5, Q=50 -> 500 real params) -python benchmark_volterra_qr.py --dataset_name APA_200MHz --pa_hidden_size 23 --model mp --K 5 --Q 50 -python benchmark_volterra_qr.py --dataset_name DPA_160MHz --pa_hidden_size 24 --model mp --K 5 --Q 50 - -# GMP (255 complex = 510 real params) -python benchmark_volterra_qr.py --dataset_name APA_200MHz --pa_hidden_size 23 --model gmp -python benchmark_volterra_qr.py --dataset_name DPA_160MHz --pa_hidden_size 24 --model gmp -``` +- Generated: `2026-07-26T20:45:35.368238+00:00` +- Git commit: `3df35e081e6e41463fa46f21778c72a748823274` +- Git branch: `benchmark-fix` +- Git working tree at launch: `dirty`; exact source hashes and start status are retained in the machine evidence. +- Python: `3.13.14` +- PyTorch: `2.13.0+cu132` +- CUDA: `13.2` +- GPU: `NVIDIA RTX PRO 6000 Blackwell Workstation Edition` +- Canonical evidence schema: `6` (published from collector schema `6`); the source archive and completed job ledger are cryptographically bound in the machine-readable evidence. +- Reproduce the full matrix: [`reproduce_benchmark_report.sh`](reproduce_benchmark_report.sh) +- Machine-readable evidence: [`results/benchmark_report_results.json`](results/benchmark_report_results.json) +- Each reproduction run writes exact commands, a verified source snapshot, raw polynomial results, checkpoints, and CSV logs to its timestamped evidence directory. diff --git a/benchmark/benchmark_results.png b/benchmark/benchmark_results.png new file mode 100644 index 0000000..9da9171 Binary files /dev/null and b/benchmark/benchmark_results.png differ diff --git a/benchmark/benchmark_training_speed.py b/benchmark/benchmark_training_speed.py deleted file mode 100644 index a065d63..0000000 --- a/benchmark/benchmark_training_speed.py +++ /dev/null @@ -1,450 +0,0 @@ -#!/usr/bin/env python3 -"""Wall-clock training-speed benchmark for OpenDPD. - -For every combination of: - role : PA (CoreModel) - DPD (CascadedModel = trainable DPD CoreModel + frozen PA CoreModel) - backbone : GRU, TCN - target #P : 200, 500, 1000 (trainable parameters) - device : cpu, cuda (if available), mps (if available) - batch size : 16, 32, 64, 128, 256 - -train one epoch on a synthetic IQ tensor and record the wall-clock time and -sample throughput. Results are rendered as rich tables in the terminal and -written to a markdown report. - -Run: - python benchmark/benchmark_training_speed.py -""" -from __future__ import annotations - -import argparse -import contextlib -import io -import json -import platform -import sys -import time -from dataclasses import asdict, dataclass -from pathlib import Path - -ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) - -import torch -import torch.nn as nn -from rich import box -from rich.console import Console -from rich.panel import Panel -from rich.table import Table -from torch.utils.data import DataLoader, TensorDataset -from tqdm import tqdm - -from models import CascadedModel, CoreModel - - -# ----------------------------------------------------------------------------- -# Configuration -# ----------------------------------------------------------------------------- -ROLES = ("PA", "DPD") -BACKBONES = ("gru", "tcn") -TARGET_PARAMS = (200, 500, 1000) -BATCH_SIZES = (16, 32, 64, 128, 256) - -# Hidden sizes chosen so the *trainable* parameter count lands close to the -# target. Exact counts (also reported in the output): -# GRU(input=2, hidden=H, output=2, num_layers=1) -> 3*H^2 + 14*H + 2 -# TCN(hidden_channels=H) -> 29*H -HIDDEN_SIZES = { - "gru": {200: 6, 500: 11, 1000: 16}, - "tcn": {200: 7, 500: 17, 1000: 34}, -} - -# DPD runs use a CascadedModel(DPD + frozen PA). The frozen PA matches the -# DPD backbone with a fixed mid-size hidden so the trainable parameter count -# corresponds purely to the DPD model. -FROZEN_PA_HIDDEN = {"gru": 11, "tcn": 17} - -INPUT_SIZE = 2 -FRAME_LENGTH = 200 -N_FRAMES = 2048 -SEED = 0 - - -# ----------------------------------------------------------------------------- -# Devices -# ----------------------------------------------------------------------------- -def discover_devices(force: list[str] | None = None) -> list[torch.device]: - if force: - return [torch.device(d) for d in force] - devices = [torch.device("cpu")] - if torch.cuda.is_available(): - devices.append(torch.device("cuda")) - mps = getattr(torch.backends, "mps", None) - if mps is not None and mps.is_available() and mps.is_built(): - devices.append(torch.device("mps")) - return devices - - -def device_short(device: torch.device) -> str: - return {"cuda": "GPU", "mps": "MPS", "cpu": "CPU"}[device.type] - - -def device_label(device: torch.device) -> str: - if device.type == "cuda": - idx = device.index if device.index is not None else 0 - try: - return f"GPU (CUDA): {torch.cuda.get_device_name(idx)}" - except Exception: - return "GPU (CUDA)" - if device.type == "mps": - return "Apple Silicon (MPS)" - return f"CPU ({platform.processor() or platform.machine()})" - - -def device_sync(device: torch.device) -> None: - if device.type == "cuda": - torch.cuda.synchronize(device) - elif device.type == "mps": - torch.mps.synchronize() - - -# ----------------------------------------------------------------------------- -# Model factory -# ----------------------------------------------------------------------------- -def build_model(role: str, backbone: str, target_p: int, - device: torch.device) -> nn.Module: - """Build a PA or DPD-cascaded model for the given config and move to device.""" - hidden = HIDDEN_SIZES[backbone][target_p] - # Suppress the "Backbone Initialized..." print from CoreModel.__init__. - with contextlib.redirect_stdout(io.StringIO()): - primary = CoreModel(input_size=INPUT_SIZE, hidden_size=hidden, - num_layers=1, backbone_type=backbone) - if role == "PA": - return primary.to(device) - pa = CoreModel(input_size=INPUT_SIZE, - hidden_size=FROZEN_PA_HIDDEN[backbone], - num_layers=1, backbone_type=backbone) - cascaded = CascadedModel(dpd_model=primary, pa_model=pa) - cascaded.freeze_pa_model() - return cascaded.to(device) - - -def trainable_params(net: nn.Module) -> int: - return sum(p.numel() for p in net.parameters() if p.requires_grad) - - -# ----------------------------------------------------------------------------- -# Single benchmark run -# ----------------------------------------------------------------------------- -@dataclass -class RunResult: - role: str - backbone: str - target_params: int - actual_params: int - device: str - batch_size: int - n_batches: int - epoch_seconds: float - samples_per_sec: float - - def to_dict(self) -> dict: - return asdict(self) - - -def run_one_epoch(role: str, backbone: str, target_p: int, - device: torch.device, batch_size: int, - features: torch.Tensor, targets: torch.Tensor, - outer_pbar: tqdm) -> RunResult: - torch.manual_seed(SEED) - net = build_model(role, backbone, target_p, device) - n_params = trainable_params(net) - - optimizer = torch.optim.AdamW( - (p for p in net.parameters() if p.requires_grad), lr=5e-4) - criterion = nn.MSELoss() - - loader = DataLoader(TensorDataset(features, targets), - batch_size=batch_size, shuffle=False, num_workers=0) - n_batches = len(loader) - net.train() - - # Warmup (one batch, not timed). Stabilises CUDA allocator / cuDNN - # autotune / first MPS kernel compile so the timed epoch is representative. - warm_x, warm_y = next(iter(loader)) - warm_x, warm_y = warm_x.to(device), warm_y.to(device) - optimizer.zero_grad() - criterion(net(warm_x), warm_y).backward() - optimizer.step() - device_sync(device) - - desc = f"{role:<3s} {backbone.upper():<3s} ~{target_p:>4d}P {device_short(device):<3s} bs={batch_size:<3d}" - inner = tqdm(loader, desc=desc, leave=False, position=1, - dynamic_ncols=True, file=sys.stdout) - - device_sync(device) - t0 = time.perf_counter() - n_samples = 0 - for x, y in inner: - x, y = x.to(device), y.to(device) - optimizer.zero_grad() - loss = criterion(net(x), y) - loss.backward() - optimizer.step() - n_samples += x.size(0) - device_sync(device) - elapsed = time.perf_counter() - t0 - inner.close() - - # Free memory before the next config. - del net, optimizer - if device.type == "cuda": - torch.cuda.empty_cache() - - outer_pbar.update(1) - return RunResult( - role=role, backbone=backbone, target_params=target_p, - actual_params=n_params, device=device_short(device), - batch_size=batch_size, n_batches=n_batches, - epoch_seconds=elapsed, - samples_per_sec=n_samples / elapsed if elapsed > 0 else float("nan"), - ) - - -# ----------------------------------------------------------------------------- -# Console rendering -# ----------------------------------------------------------------------------- -def _row_keys(results: list[RunResult]): - return sorted({(r.role, r.backbone, r.target_params, r.actual_params) - for r in results}) - - -def _find(results: list[RunResult], role: str, backbone: str, - target_p: int, batch_size: int) -> RunResult | None: - for r in results: - if (r.role == role and r.backbone == backbone - and r.target_params == target_p and r.batch_size == batch_size): - return r - return None - - -def make_results_table(results: list[RunResult], device_name: str) -> Table: - table = Table( - title=(f"[bold cyan]{device_name}[/bold cyan] " - f"[dim]— seconds / epoch • samples / sec[/dim]"), - box=box.SIMPLE_HEAVY, - title_style="bold", - header_style="bold magenta", - pad_edge=False, - ) - table.add_column("Role", justify="center", style="bold cyan") - table.add_column("Backbone", justify="center", style="bold") - table.add_column("Target #P", justify="right", style="dim") - table.add_column("Actual #P", justify="right") - for bs in BATCH_SIZES: - table.add_column(f"bs={bs}", justify="right") - - last_role: str | None = None - for role, backbone, tgt, act in _row_keys(results): - if last_role is not None and role != last_role: - table.add_section() - last_role = role - cells = [role, backbone.upper(), str(tgt), f"{act:,}"] - for bs in BATCH_SIZES: - r = _find(results, role, backbone, tgt, bs) - if r is None: - cells.append("[red]—[/red]") - else: - cells.append(f"[bold]{r.epoch_seconds:.3f}s[/bold]\n" - f"[green]{r.samples_per_sec:,.0f}[/green]") - table.add_row(*cells) - return table - - -def render_console(console: Console, all_results: list[RunResult], - device_descs: dict[str, str]) -> None: - by_device: dict[str, list[RunResult]] = {} - for r in all_results: - by_device.setdefault(r.device, []).append(r) - for device_short_name, rs in by_device.items(): - console.print() - console.print(make_results_table(rs, device_descs.get( - device_short_name, device_short_name))) - - -# ----------------------------------------------------------------------------- -# Markdown report -# ----------------------------------------------------------------------------- -def render_markdown(all_results: list[RunResult], - device_descs: dict[str, str], - sys_info: dict[str, str], total_seconds: float, - output: Path) -> None: - lines: list[str] = [] - lines.append("# OpenDPD Training-Speed Benchmark") - lines.append("") - lines.append( - "Wall-clock benchmark of one training epoch for **PA** and **DPD** " - "models with **GRU** and **TCN** backbones at three parameter scales " - "(~200, ~500, ~1000), across batch sizes " - f"{{{', '.join(str(b) for b in BATCH_SIZES)}}} on each available device." - ) - lines.append("") - - # System info - lines.append("## System") - lines.append("") - lines.append("| Field | Value |") - lines.append("|---|---|") - for k, v in sys_info.items(): - lines.append(f"| {k} | {v} |") - lines.append(f"| Total benchmark wall-clock | {total_seconds:.1f} s |") - lines.append("") - - # Configuration - lines.append("## Configuration") - lines.append("") - lines.append(f"- Synthetic dataset: **{N_FRAMES}** IQ frames of length " - f"**{FRAME_LENGTH}** (random, seed={SEED})") - lines.append("- Optimiser: **AdamW** (lr=5e-4), MSE loss") - lines.append("- 1 warmup batch (untimed) + 1 full epoch (timed) per configuration") - lines.append( - "- DPD runs use **CascadedModel(DPD + frozen PA)**; the frozen PA " - f"uses the same backbone as the DPD with hidden_size = " - f"{FROZEN_PA_HIDDEN['gru']} (GRU) / {FROZEN_PA_HIDDEN['tcn']} (TCN)." - ) - lines.append("") - lines.append("**Hidden sizes used to hit each parameter target**") - lines.append("") - lines.append("| Backbone | ~200 P | ~500 P | ~1000 P |") - lines.append("|---|---|---|---|") - for bb in BACKBONES: - cells = [bb.upper()] + [str(HIDDEN_SIZES[bb][p]) for p in TARGET_PARAMS] - lines.append("| " + " | ".join(cells) + " |") - lines.append("") - - # Results - lines.append("## Results") - lines.append("") - by_device: dict[str, list[RunResult]] = {} - for r in all_results: - by_device.setdefault(r.device, []).append(r) - - for device_short_name, rs in by_device.items(): - lines.append(f"### {device_descs.get(device_short_name, device_short_name)}") - lines.append("") - header = (["Role", "Backbone", "Target #P", "Actual #P"] - + [f"bs={bs}" for bs in BATCH_SIZES]) - - lines.append("**Wall-clock per epoch (seconds)**") - lines.append("") - lines.append("| " + " | ".join(header) + " |") - lines.append("|" + "|".join(["---"] * len(header)) + "|") - for role, backbone, tgt, act in _row_keys(rs): - row = [role, backbone.upper(), str(tgt), f"{act:,}"] - for bs in BATCH_SIZES: - r = _find(rs, role, backbone, tgt, bs) - row.append(f"{r.epoch_seconds:.3f}" if r else "—") - lines.append("| " + " | ".join(row) + " |") - lines.append("") - - lines.append("**Throughput (samples / sec)**") - lines.append("") - lines.append("| " + " | ".join(header) + " |") - lines.append("|" + "|".join(["---"] * len(header)) + "|") - for role, backbone, tgt, act in _row_keys(rs): - row = [role, backbone.upper(), str(tgt), f"{act:,}"] - for bs in BATCH_SIZES: - r = _find(rs, role, backbone, tgt, bs) - row.append(f"{r.samples_per_sec:,.0f}" if r else "—") - lines.append("| " + " | ".join(row) + " |") - lines.append("") - - output.write_text("\n".join(lines) + "\n") - - -# ----------------------------------------------------------------------------- -# Main -# ----------------------------------------------------------------------------- -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--devices", nargs="*", default=None, - choices=["cpu", "cuda", "mps"], - help="Override which devices to benchmark (default: auto).") - p.add_argument("--output", type=Path, - default=Path(__file__).resolve().parent - / "benchmark_training_speed_report.md", - help="Markdown report output path.") - p.add_argument("--json-output", type=Path, default=None, - help="Optional path to also dump raw results as JSON.") - return p.parse_args() - - -def main() -> None: - args = parse_args() - console = Console() - - devices = discover_devices(args.devices) - device_descs = {device_short(d): device_label(d) for d in devices} - - g = torch.Generator().manual_seed(SEED) - features = torch.randn(N_FRAMES, FRAME_LENGTH, INPUT_SIZE, generator=g) - targets = torch.randn(N_FRAMES, FRAME_LENGTH, INPUT_SIZE, generator=g) - - sys_info = { - "Platform": f"{platform.system()} {platform.release()} ({platform.machine()})", - "Python": platform.python_version(), - "PyTorch": torch.__version__, - "Devices": " | ".join(device_descs.values()), - "CPU threads (torch)": str(torch.get_num_threads()), - } - panel_text = "\n".join(f"[bold]{k:<22}[/bold] {v}" for k, v in sys_info.items()) - console.print(Panel( - panel_text, - title="[bold cyan]OpenDPD Training-Speed Benchmark[/bold cyan]", - border_style="cyan", expand=False, - )) - - configs = [ - (role, backbone, tp, dev, bs) - for dev in devices - for role in ROLES - for backbone in BACKBONES - for tp in TARGET_PARAMS - for bs in BATCH_SIZES - ] - console.print(f"[dim]Total configurations: {len(configs)}[/dim]\n") - - results: list[RunResult] = [] - bench_t0 = time.perf_counter() - - overall = tqdm(total=len(configs), desc="overall", position=0, - leave=True, dynamic_ncols=True, file=sys.stdout) - try: - for role, backbone, tp, dev, bs in configs: - try: - results.append(run_one_epoch( - role, backbone, tp, dev, bs, features, targets, overall)) - except RuntimeError as e: - overall.update(1) - console.print( - f"[red]Failed[/red] {role}/{backbone}/p{tp}/" - f"{device_short(dev)}/bs{bs}: {e}") - finally: - overall.close() - - total_seconds = time.perf_counter() - bench_t0 - - render_console(console, results, device_descs) - render_markdown(results, device_descs, sys_info, total_seconds, args.output) - console.print(f"\n[green]Markdown report written to:[/green] {args.output}") - - if args.json_output: - args.json_output.write_text( - json.dumps([r.to_dict() for r in results], indent=2)) - console.print(f"[green]JSON results written to:[/green] {args.json_output}") - - -if __name__ == "__main__": - main() diff --git a/benchmark/benchmark_training_speed_report.md b/benchmark/benchmark_training_speed_report.md deleted file mode 100644 index 815a46d..0000000 --- a/benchmark/benchmark_training_speed_report.md +++ /dev/null @@ -1,103 +0,0 @@ -# OpenDPD Training-Speed Benchmark - -Wall-clock benchmark of one training epoch for **PA** and **DPD** models with **GRU** and **TCN** backbones at three parameter scales (~200, ~500, ~1000), across batch sizes {16, 32, 64, 128, 256} on each available device. - -## System - -| Field | Value | -|---|---| -| Platform | Linux 6.17.0-19-generic (x86_64) | -| Python | 3.13.12 | -| PyTorch | 2.11.0+cu128 | -| Devices | CPU (x86_64) | GPU (CUDA): NVIDIA GeForce RTX 4090 | -| CPU threads (torch) | 16 | -| Total benchmark wall-clock | 69.5 s | - -## Configuration - -- Synthetic dataset: **2048** IQ frames of length **200** (random, seed=0) -- Optimiser: **AdamW** (lr=5e-4), MSE loss -- 1 warmup batch (untimed) + 1 full epoch (timed) per configuration -- DPD runs use **CascadedModel(DPD + frozen PA)**; the frozen PA uses the same backbone as the DPD with hidden_size = 11 (GRU) / 17 (TCN). - -**Hidden sizes used to hit each parameter target** - -| Backbone | ~200 P | ~500 P | ~1000 P | -|---|---|---|---| -| GRU | 6 | 11 | 16 | -| TCN | 7 | 17 | 34 | - -## Results - -### CPU (x86_64) - -**Wall-clock per epoch (seconds)** - -| Role | Backbone | Target #P | Actual #P | bs=16 | bs=32 | bs=64 | bs=128 | bs=256 | -|---|---|---|---|---|---|---|---|---| -| DPD | GRU | 200 | 194 | 5.343 | 2.885 | 1.666 | 0.969 | 0.682 | -| DPD | GRU | 500 | 519 | 5.162 | 2.702 | 1.807 | 1.033 | 0.671 | -| DPD | GRU | 1000 | 994 | 5.038 | 2.888 | 1.734 | 1.040 | 0.658 | -| DPD | TCN | 200 | 203 | 0.510 | 0.263 | 0.204 | 0.154 | 0.153 | -| DPD | TCN | 500 | 493 | 0.486 | 0.389 | 0.243 | 0.214 | 0.206 | -| DPD | TCN | 1000 | 986 | 0.559 | 0.439 | 0.333 | 0.316 | 0.317 | -| PA | GRU | 200 | 194 | 2.874 | 1.609 | 0.773 | 0.562 | 0.452 | -| PA | GRU | 500 | 519 | 3.997 | 2.123 | 0.853 | 0.768 | 0.391 | -| PA | GRU | 1000 | 994 | 2.806 | 1.478 | 0.884 | 0.539 | 0.361 | -| PA | TCN | 200 | 203 | 0.239 | 0.143 | 0.122 | 0.071 | 0.063 | -| PA | TCN | 500 | 493 | 0.294 | 0.231 | 0.153 | 0.122 | 0.135 | -| PA | TCN | 1000 | 986 | 0.402 | 0.288 | 0.228 | 0.216 | 0.456 | - -**Throughput (samples / sec)** - -| Role | Backbone | Target #P | Actual #P | bs=16 | bs=32 | bs=64 | bs=128 | bs=256 | -|---|---|---|---|---|---|---|---|---| -| DPD | GRU | 200 | 194 | 383 | 710 | 1,229 | 2,114 | 3,005 | -| DPD | GRU | 500 | 519 | 397 | 758 | 1,133 | 1,983 | 3,051 | -| DPD | GRU | 1000 | 994 | 406 | 709 | 1,181 | 1,969 | 3,111 | -| DPD | TCN | 200 | 203 | 4,020 | 7,791 | 10,040 | 13,284 | 13,390 | -| DPD | TCN | 500 | 493 | 4,218 | 5,268 | 8,433 | 9,562 | 9,960 | -| DPD | TCN | 1000 | 986 | 3,664 | 4,662 | 6,142 | 6,473 | 6,460 | -| PA | GRU | 200 | 194 | 713 | 1,273 | 2,651 | 3,642 | 4,527 | -| PA | GRU | 500 | 519 | 512 | 965 | 2,401 | 2,667 | 5,237 | -| PA | GRU | 1000 | 994 | 730 | 1,385 | 2,317 | 3,799 | 5,671 | -| PA | TCN | 200 | 203 | 8,585 | 14,331 | 16,735 | 28,660 | 32,636 | -| PA | TCN | 500 | 493 | 6,970 | 8,860 | 13,428 | 16,797 | 15,206 | -| PA | TCN | 1000 | 986 | 5,092 | 7,111 | 8,991 | 9,473 | 4,493 | - -### GPU (CUDA): NVIDIA GeForce RTX 4090 - -**Wall-clock per epoch (seconds)** - -| Role | Backbone | Target #P | Actual #P | bs=16 | bs=32 | bs=64 | bs=128 | bs=256 | -|---|---|---|---|---|---|---|---|---| -| DPD | GRU | 200 | 194 | 0.130 | 0.068 | 0.038 | 0.022 | 0.014 | -| DPD | GRU | 500 | 519 | 0.129 | 0.068 | 0.039 | 0.023 | 0.015 | -| DPD | GRU | 1000 | 994 | 0.129 | 0.068 | 0.039 | 0.023 | 0.015 | -| DPD | TCN | 200 | 203 | 0.169 | 0.087 | 0.047 | 0.026 | 0.016 | -| DPD | TCN | 500 | 493 | 0.168 | 0.087 | 0.047 | 0.026 | 0.016 | -| DPD | TCN | 1000 | 986 | 0.182 | 0.108 | 0.057 | 0.031 | 0.019 | -| PA | GRU | 200 | 194 | 0.109 | 0.049 | 0.026 | 0.017 | 0.011 | -| PA | GRU | 500 | 519 | 0.092 | 0.049 | 0.027 | 0.016 | 0.011 | -| PA | GRU | 1000 | 994 | 0.096 | 0.060 | 0.029 | 0.018 | 0.012 | -| PA | TCN | 200 | 203 | 0.096 | 0.051 | 0.028 | 0.017 | 0.013 | -| PA | TCN | 500 | 493 | 0.096 | 0.051 | 0.029 | 0.017 | 0.012 | -| PA | TCN | 1000 | 986 | 0.096 | 0.052 | 0.029 | 0.017 | 0.012 | - -**Throughput (samples / sec)** - -| Role | Backbone | Target #P | Actual #P | bs=16 | bs=32 | bs=64 | bs=128 | bs=256 | -|---|---|---|---|---|---|---|---|---| -| DPD | GRU | 200 | 194 | 15,764 | 30,061 | 54,151 | 91,706 | 141,334 | -| DPD | GRU | 500 | 519 | 15,834 | 29,995 | 51,872 | 90,196 | 140,680 | -| DPD | GRU | 1000 | 994 | 15,829 | 30,161 | 52,875 | 90,122 | 139,074 | -| DPD | TCN | 200 | 203 | 12,119 | 23,507 | 44,036 | 77,912 | 127,690 | -| DPD | TCN | 500 | 493 | 12,167 | 23,432 | 43,802 | 77,883 | 127,460 | -| DPD | TCN | 1000 | 986 | 11,239 | 19,006 | 36,172 | 66,221 | 109,519 | -| PA | GRU | 200 | 194 | 18,857 | 41,602 | 79,253 | 122,231 | 185,570 | -| PA | GRU | 500 | 519 | 22,209 | 42,141 | 76,839 | 128,057 | 188,255 | -| PA | GRU | 1000 | 994 | 21,411 | 33,889 | 69,896 | 115,885 | 173,196 | -| PA | TCN | 200 | 203 | 21,419 | 40,277 | 72,349 | 117,288 | 160,110 | -| PA | TCN | 500 | 493 | 21,270 | 39,913 | 71,851 | 119,467 | 176,378 | -| PA | TCN | 1000 | 986 | 21,269 | 39,686 | 71,682 | 118,658 | 175,314 | - diff --git a/benchmark/benchmark_volterra.py b/benchmark/benchmark_volterra.py new file mode 100644 index 0000000..cf55d3c --- /dev/null +++ b/benchmark/benchmark_volterra.py @@ -0,0 +1,1025 @@ +"""Polynomial PA-modeling and indirect-learning DPD benchmarks. + +The benchmark supports two complex polynomial structures: + +* MP: a memory polynomial. +* GMP: a generalized memory polynomial with aligned, lagging, and leading + envelope terms. + +For PA modeling, the complex coefficients are identified directly from +``Phi(x) w ~= y``. For DPD, an indirect learning architecture (ILA) first +identifies the postdistorter from ``Phi(y / G) w ~= x``, copies those +coefficients to the predistorter, and evaluates the result through a trained +PA model. + +Large benchmark systems are solved on CUDA after L2 column scaling. The +well-conditioned MP and ILA systems use ``torch.linalg.lstsq`` with ``gels``. +GMP PA modeling can instead use a rank-controlled truncated SVD, because its +correlated cross-term basis does not satisfy CUDA ``gels``'s full-rank +assumption on every dataset. +""" + +from __future__ import annotations + +import argparse +from functools import partial +import hashlib +import json +import math +import os +from pathlib import Path +import time +from typing import Callable + +import numpy as np +import torch + +import models as model +from modules.data_collector import load_dataset +from utils import metrics +from utils.util import count_net_params, set_target_gain + + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +# --------------------------------------------------------------------------- +# NumPy basis implementations (small inputs, reference tests, and API users) +# --------------------------------------------------------------------------- + +def build_mp_basis(x_complex, K, Q): + """Build an MP basis with columns ``x(n-q) * |x(n-q)|**k``.""" + _validate_positive_configuration({"K": K, "Q": Q}) + x_complex = np.asarray(x_complex, dtype=np.complex128) + n_samples = len(x_complex) + phi = np.zeros((n_samples, K * Q), dtype=np.complex128) + column = 0 + for k in range(K): + for q in range(Q): + delayed = _delay(x_complex, q) + phi[:, column] = delayed * np.abs(delayed) ** k + column += 1 + return phi + + +def _delay(x, delay): + """Shift a one-dimensional complex signal and zero-fill the boundary.""" + x = np.asarray(x) + n_samples = len(x) + output = np.zeros(n_samples, dtype=x.dtype) + if 0 <= delay < n_samples: + output[delay:] = x[:n_samples - delay] + elif -n_samples < delay < 0: + output[:n_samples + delay] = x[-delay:] + return output + + +def build_gmp_basis(x_complex, Ka, La, Kb, Lb, Mb, Kc, Lc, Mc): + """Build the aligned, lagging, and leading GMP basis.""" + configuration = { + "Ka": Ka, + "La": La, + "Kb": Kb, + "Lb": Lb, + "Mb": Mb, + "Kc": Kc, + "Lc": Lc, + "Mc": Mc, + } + _validate_positive_configuration(configuration) + x_complex = np.asarray(x_complex, dtype=np.complex128) + n_samples = len(x_complex) + n_columns = gmp_coefficient_count(configuration) + phi = np.zeros((n_samples, n_columns), dtype=np.complex128) + column = 0 + + for k in range(Ka): + for q in range(La): + xq = _delay(x_complex, q) + phi[:, column] = xq * np.abs(xq) ** k + column += 1 + + for k in range(1, Kb + 1): + for q in range(Lb): + for lag in range(1, Mb + 1): + xq = _delay(x_complex, q) + envelope = _delay(x_complex, q + lag) + phi[:, column] = xq * np.abs(envelope) ** k + column += 1 + + for k in range(1, Kc + 1): + for q in range(Lc): + for lead in range(1, Mc + 1): + xq = _delay(x_complex, q) + envelope = _delay(x_complex, q - lead) + phi[:, column] = xq * np.abs(envelope) ** k + column += 1 + + return phi + + +def build_segmented_numpy_basis( + x_complex: np.ndarray, + basis_builder: Callable[[np.ndarray], np.ndarray], + segment_length: int, +) -> np.ndarray: + """Build a reference basis while resetting delays at every segment.""" + if segment_length <= 0: + raise ValueError("segment_length must be positive") + x_complex = np.asarray(x_complex) + blocks = [ + basis_builder(x_complex[start:start + segment_length]) + for start in range(0, len(x_complex), segment_length) + ] + if not blocks: + raise ValueError("cannot build a basis for an empty signal") + return np.concatenate(blocks, axis=0) + + +# --------------------------------------------------------------------------- +# Memory-conscious Torch basis implementations used by the benchmark +# --------------------------------------------------------------------------- + +def _torch_delay(x: torch.Tensor, delay: int) -> torch.Tensor: + output = torch.zeros_like(x) + n_samples = x.numel() + if 0 <= delay < n_samples: + output[delay:] = x[:n_samples - delay] + elif -n_samples < delay < 0: + output[:n_samples + delay] = x[-delay:] + return output + + +def _torch_delay_columns(x: torch.Tensor, delays: list[int]) -> torch.Tensor: + return torch.stack([_torch_delay(x, delay) for delay in delays], dim=1) + + +def _build_mp_basis_torch_block( + x: torch.Tensor, + configuration: dict[str, int], +) -> torch.Tensor: + k_count = configuration["K"] + memory_depth = configuration["Q"] + delayed = _torch_delay_columns(x, list(range(memory_depth))) + magnitudes = torch.abs(delayed) + phi = torch.empty( + (x.numel(), k_count * memory_depth), + dtype=x.dtype, + device=x.device, + ) + for k in range(k_count): + start = k * memory_depth + phi[:, start:start + memory_depth] = delayed * magnitudes.pow(k) + return phi + + +def _build_gmp_basis_torch_block( + x: torch.Tensor, + configuration: dict[str, int], +) -> torch.Tensor: + Ka, La = configuration["Ka"], configuration["La"] + Kb, Lb, Mb = configuration["Kb"], configuration["Lb"], configuration["Mb"] + Kc, Lc, Mc = configuration["Kc"], configuration["Lc"], configuration["Mc"] + phi = torch.empty( + (x.numel(), gmp_coefficient_count(configuration)), + dtype=x.dtype, + device=x.device, + ) + column = 0 + + aligned = _torch_delay_columns(x, list(range(La))) + aligned_magnitude = torch.abs(aligned) + for k in range(Ka): + phi[:, column:column + La] = aligned * aligned_magnitude.pow(k) + column += La + + lag_base = _torch_delay_columns(x, list(range(Lb))) + lag_delays = [q + lag for q in range(Lb) for lag in range(1, Mb + 1)] + lag_envelopes = torch.abs(_torch_delay_columns(x, lag_delays)).reshape( + x.numel(), Lb, Mb + ) + for k in range(1, Kb + 1): + block = lag_base.unsqueeze(-1) * lag_envelopes.pow(k) + width = Lb * Mb + phi[:, column:column + width] = block.reshape(x.numel(), width) + column += width + + lead_base = _torch_delay_columns(x, list(range(Lc))) + lead_delays = [q - lead for q in range(Lc) for lead in range(1, Mc + 1)] + lead_envelopes = torch.abs(_torch_delay_columns(x, lead_delays)).reshape( + x.numel(), Lc, Mc + ) + for k in range(1, Kc + 1): + block = lead_base.unsqueeze(-1) * lead_envelopes.pow(k) + width = Lc * Mc + phi[:, column:column + width] = block.reshape(x.numel(), width) + column += width + + if column != phi.shape[1]: + raise RuntimeError("internal GMP column-count mismatch") + return phi + + +@torch.inference_mode() +def build_torch_basis( + x: torch.Tensor, + polynomial_model: str, + configuration: dict[str, int], + segment_length: int, +) -> torch.Tensor: + """Build a design matrix on ``x.device`` with per-segment delay resets.""" + if x.ndim != 1 or not torch.is_complex(x): + raise ValueError("x must be a one-dimensional complex Torch tensor") + if x.numel() == 0: + raise ValueError("cannot build a basis for an empty signal") + if segment_length <= 0: + raise ValueError("segment_length must be positive") + _validate_positive_configuration(configuration) + + n_columns = coefficient_count(polynomial_model, configuration) + phi = torch.empty( + (x.numel(), n_columns), + dtype=x.dtype, + device=x.device, + ) + build_block = ( + _build_mp_basis_torch_block + if polynomial_model == "mp" + else _build_gmp_basis_torch_block + ) + for start in range(0, x.numel(), segment_length): + stop = min(start + segment_length, x.numel()) + phi[start:stop] = build_block(x[start:stop], configuration) + return phi + + +# --------------------------------------------------------------------------- +# Configuration and least-squares helpers +# --------------------------------------------------------------------------- + +def _validate_positive_configuration(configuration: dict[str, int]) -> None: + for name, value in configuration.items(): + if not isinstance(value, (int, np.integer)) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + + +def gmp_coefficient_count(configuration: dict[str, int]) -> int: + return ( + configuration["Ka"] * configuration["La"] + + configuration["Kb"] * configuration["Lb"] * configuration["Mb"] + + configuration["Kc"] * configuration["Lc"] * configuration["Mc"] + ) + + +def coefficient_count( + polynomial_model: str, + configuration: dict[str, int], +) -> int: + if polynomial_model == "mp": + return configuration["K"] * configuration["Q"] + if polynomial_model == "gmp": + return gmp_coefficient_count(configuration) + raise ValueError(f"unsupported polynomial model: {polynomial_model!r}") + + +def basis_configuration(args: argparse.Namespace) -> dict[str, int]: + if args.model == "mp": + return {"K": args.K, "Q": args.Q} + return { + "Ka": args.Ka, + "La": args.La, + "Kb": args.Kb, + "Lb": args.Lb, + "Mb": args.Mb, + "Kc": args.Kc, + "Lc": args.Lc, + "Mc": args.Mc, + } + + +def select_basis(args): + """Return the NumPy reference builder, coefficient count, and banner.""" + configuration = basis_configuration(args) + if args.model == "gmp": + builder = partial(build_gmp_basis, **configuration) + banner = ( + "Polynomial GMP " + + ", ".join(f"{name}={value}" for name, value in configuration.items()) + ) + else: + builder = partial(build_mp_basis, **configuration) + banner = f"Polynomial MP K={args.K}, Q={args.Q}" + return builder, coefficient_count(args.model, configuration), banner + + +def resolve_solver_device(requested: str, cuda_index: int) -> torch.device: + requested = requested.strip().lower() + if requested == "cuda": + requested = f"cuda:{cuda_index}" + device = torch.device(requested) + if device.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested for the least-squares solve but is unavailable") + if device.index is None: + device = torch.device(f"cuda:{cuda_index}") + elif device.type != "cpu": + raise ValueError("solver device must be 'cpu', 'cuda', or 'cuda:INDEX'") + return device + + +def resolve_complex_dtype(requested: str) -> torch.dtype: + mapping = { + "complex64": torch.complex64, + "torch.complex64": torch.complex64, + "complex128": torch.complex128, + "torch.complex128": torch.complex128, + } + try: + return mapping[requested.lower()] + except KeyError as error: + raise ValueError("solver dtype must be complex64 or complex128") from error + + +def _synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +@torch.inference_mode() +def fit_complex_least_squares( + basis_input: np.ndarray, + target: np.ndarray, + *, + polynomial_model: str, + configuration: dict[str, int], + segment_length: int, + device: torch.device, + dtype: torch.dtype, + solver_mode: str = "gels", + svd_rcond: float | None = None, +) -> tuple[torch.Tensor, dict[str, float | int | str | None]]: + """Fit complex coefficients after L2 column scaling.""" + basis_input = np.asarray(basis_input) + target = np.asarray(target) + if basis_input.ndim != 1 or target.ndim != 1: + raise ValueError("least-squares input and target must be one-dimensional") + if basis_input.shape != target.shape: + raise ValueError("least-squares input and target must have the same shape") + if not np.all(np.isfinite(basis_input)) or not np.all(np.isfinite(target)): + raise ValueError("least-squares input and target must contain only finite values") + + n_columns = coefficient_count(polynomial_model, configuration) + if basis_input.size < n_columns: + raise ValueError( + f"least-squares system is underdetermined: " + f"{basis_input.size} observations for {n_columns} coefficients" + ) + if solver_mode not in {"gels", "truncated_svd"}: + raise ValueError(f"unsupported solver mode: {solver_mode!r}") + if solver_mode == "gels" and svd_rcond is not None: + raise ValueError("svd_rcond is only valid for truncated_svd") + if solver_mode == "truncated_svd": + if svd_rcond is None or not math.isfinite(svd_rcond): + raise ValueError( + "truncated_svd requires a finite svd_rcond" + ) + if not 0.0 < svd_rcond < 1.0: + raise ValueError("svd_rcond must be between zero and one") + + _synchronize(device) + started = time.perf_counter() + x_tensor = torch.as_tensor(basis_input, dtype=dtype, device=device) + target_tensor = torch.as_tensor(target, dtype=dtype, device=device) + phi = build_torch_basis( + x_tensor, + polynomial_model, + configuration, + segment_length, + ) + _synchronize(device) + basis_seconds = time.perf_counter() - started + + column_norms = torch.linalg.vector_norm(phi, dim=0) + if not bool(torch.all(torch.isfinite(column_norms)).item()): + raise RuntimeError("basis column norms are not finite") + if bool(torch.any(column_norms == 0).item()): + zero_count = int(torch.count_nonzero(column_norms == 0).item()) + raise RuntimeError(f"basis contains {zero_count} all-zero columns") + smallest_norm = float(column_norms.min().item()) + largest_norm = float(column_norms.max().item()) + phi.div_(column_norms) + + _synchronize(device) + solve_started = time.perf_counter() + least_squares_rank: int | None = None + singular_value_max: float | None = None + singular_value_min: float | None = None + singular_value_cutoff: float | None = None + retained_singular_value_min: float | None = None + condition_number: float | None = None + if solver_mode == "gels": + solver_implementation = "torch.linalg.lstsq" + solver_driver = "gels" + regularization = None + try: + solution_scaled = torch.linalg.lstsq( + phi, + target_tensor.unsqueeze(1), + driver=solver_driver, + ).solution.squeeze(1) + except RuntimeError as error: + raise RuntimeError( + "torch.linalg.lstsq failed for the polynomial benchmark. " + "The CUDA gels driver assumes a full-rank design matrix; " + "verify the basis configuration and available accelerator " + "memory." + ) from error + else: + solver_implementation = "torch.linalg.svd" + solver_driver = "gesvdj" if device.type == "cuda" else "default" + regularization = "truncated_svd" + svd_driver = "gesvdj" if device.type == "cuda" else None + try: + left, singular_values, right_h = torch.linalg.svd( + phi, + full_matrices=False, + driver=svd_driver, + ) + except RuntimeError as error: + raise RuntimeError( + "torch.linalg.svd failed for the polynomial benchmark; " + "verify the basis configuration and available accelerator " + "memory." + ) from error + singular_value_max = float(singular_values[0].item()) + singular_value_min = float(singular_values[-1].item()) + singular_value_cutoff = singular_value_max * float(svd_rcond) + retained = singular_values > singular_value_cutoff + least_squares_rank = int(torch.count_nonzero(retained).item()) + if least_squares_rank == 0: + raise RuntimeError( + "truncated SVD rejected every singular direction" + ) + retained_singular_value_min = float( + singular_values[retained][-1].item() + ) + condition_number = ( + singular_value_max / singular_value_min + if singular_value_min > 0.0 + else None + ) + projected_target = left[:, retained].mH @ target_tensor + solution_scaled = right_h[retained].mH @ ( + projected_target / singular_values[retained] + ) + _synchronize(device) + solve_seconds = time.perf_counter() - solve_started + + fitted = phi @ solution_scaled + residual_norm = torch.linalg.vector_norm(fitted - target_tensor) + target_norm = torch.linalg.vector_norm(target_tensor) + relative_residual = residual_norm / target_norm + coefficients = solution_scaled / column_norms + + if not bool(torch.all(torch.isfinite(coefficients)).item()): + raise RuntimeError("least-squares coefficients are not finite") + relative_residual_value = float(relative_residual.item()) + if not math.isfinite(relative_residual_value): + raise RuntimeError("least-squares training relative residual is not finite") + coefficient_l2_norm = float(torch.linalg.vector_norm(coefficients).item()) + if not math.isfinite(coefficient_l2_norm): + raise RuntimeError("least-squares coefficient norm is not finite") + + diagnostics: dict[str, float | int | str | None] = { + "observations": int(basis_input.size), + "columns": int(n_columns), + # The gels driver does not estimate numerical rank. Its documented + # contract assumes that the design matrix is full rank. + "least_squares_rank": least_squares_rank, + "solver_mode": solver_mode, + "solver_implementation": solver_implementation, + "solver_driver": solver_driver, + "regularization": regularization, + "svd_rcond": svd_rcond, + "singular_value_max": singular_value_max, + "singular_value_min": singular_value_min, + "singular_value_cutoff": singular_value_cutoff, + "retained_singular_value_min": retained_singular_value_min, + "condition_number": condition_number, + "coefficient_l2_norm": coefficient_l2_norm, + "training_relative_residual": relative_residual_value, + "training_residual_nmse_db": float( + 20.0 * math.log10(max(relative_residual_value, np.finfo(float).tiny)) + ), + "basis_build_seconds": float(basis_seconds), + "solve_seconds": float(solve_seconds), + "column_norm_min": smallest_norm, + "column_norm_max": largest_norm, + } + return coefficients, diagnostics + + +@torch.inference_mode() +def apply_polynomial( + input_complex: np.ndarray, + coefficients: torch.Tensor, + *, + polynomial_model: str, + configuration: dict[str, int], + segment_length: int, +) -> np.ndarray: + """Apply fitted polynomial coefficients with the training boundary policy.""" + input_tensor = torch.as_tensor( + np.asarray(input_complex), + dtype=coefficients.dtype, + device=coefficients.device, + ) + phi = build_torch_basis( + input_tensor, + polynomial_model, + configuration, + segment_length, + ) + prediction = phi @ coefficients + if not bool(torch.all(torch.isfinite(prediction)).item()): + raise RuntimeError("polynomial prediction contains non-finite values") + return prediction.cpu().numpy() + + +# --------------------------------------------------------------------------- +# Data, PA, and metric helpers +# --------------------------------------------------------------------------- + +def iq_to_complex(iq): + iq = np.asarray(iq) + return iq[:, 0] + 1j * iq[:, 1] + + +def complex_to_iq(complex_signal): + complex_signal = np.asarray(complex_signal) + return np.stack( + [complex_signal.real, complex_signal.imag], + axis=-1, + ).astype(np.float32) + + +def sha256_file(path): + digest = hashlib.sha256() + with open(path, "rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def pa_model_id( + *, + seed: int, + backbone: str, + hidden_size: int, + frame_length: int, + parameter_count: int, +) -> str: + return ( + f"PA_S_{seed}_M_{backbone.upper()}_H_{hidden_size}_" + f"F_{frame_length}_P_{parameter_count}" + ) + + +def resolve_pa_checkpoint( + args: argparse.Namespace, + parameter_count: int, +) -> tuple[Path, str]: + model_id = pa_model_id( + seed=args.seed, + backbone=args.pa_backbone, + hidden_size=args.pa_hidden_size, + frame_length=args.frame_length, + parameter_count=parameter_count, + ) + if args.pa_checkpoint: + checkpoint = Path(args.pa_checkpoint).expanduser() + if not checkpoint.is_absolute(): + checkpoint = REPO_ROOT / checkpoint + else: + checkpoint = ( + REPO_ROOT + / "save" + / args.dataset_name + / "train_pa" + / f"{model_id}.pt" + ) + checkpoint = checkpoint.resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"PA checkpoint does not exist: {checkpoint}") + return checkpoint, model_id + + +def _load_state_dict(path: Path): + try: + state = torch.load(path, map_location="cpu", weights_only=True) + except TypeError: # Compatibility with older supported Torch releases. + state = torch.load(path, map_location="cpu") + if isinstance(state, dict) and "state_dict" in state: + state = state["state_dict"] + if not isinstance(state, dict): + raise TypeError(f"PA checkpoint is not a state dictionary: {path}") + return state + + +def load_pa_model( + args: argparse.Namespace, +) -> tuple[torch.nn.Module, torch.device, Path, str, int]: + if args.device < 0: + pa_device = torch.device("cpu") + else: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA PA evaluation was requested but CUDA is unavailable") + pa_device = torch.device(f"cuda:{args.device}") + + pa_net = model.CoreModel( + input_size=2, + hidden_size=args.pa_hidden_size, + num_layers=args.pa_num_layers, + backbone_type=args.pa_backbone, + ) + parameter_count = count_net_params(pa_net) + checkpoint, model_id = resolve_pa_checkpoint(args, parameter_count) + pa_net.load_state_dict(_load_state_dict(checkpoint), strict=True) + pa_net = pa_net.to(pa_device) + pa_net.eval() + return pa_net, pa_device, checkpoint, model_id, parameter_count + + +@torch.inference_mode() +def run_pa_model(iq_signal, pa_net, device, nperseg): + """Evaluate the PA with one independent recurrent state per segment.""" + iq_signal = np.asarray(iq_signal, dtype=np.float32) + n_samples = len(iq_signal) + segments = [] + for start in range(0, n_samples, nperseg): + segment = iq_signal[start:start + nperseg] + if len(segment) < nperseg: + segment = np.vstack( + [ + segment, + np.zeros((nperseg - len(segment), 2), dtype=np.float32), + ] + ) + segments.append(segment) + if not segments: + raise ValueError("cannot evaluate an empty signal") + tensor = torch.as_tensor(np.stack(segments), dtype=torch.float32, device=device) + output = pa_net(tensor).cpu().numpy() + return output.reshape(-1, 2)[:n_samples] + + +def compute_metrics(prediction_iq, ground_truth_iq, spec): + """Compute metrics with the same segmentation and definitions as training.""" + nperseg = int(spec["nperseg"]) + n_samples = min(len(prediction_iq), len(ground_truth_iq)) + n_segments = n_samples // nperseg + if n_segments == 0: + raise ValueError( + f"metric input needs at least one complete {nperseg}-sample segment" + ) + usable = n_segments * nperseg + prediction = np.asarray(prediction_iq[:usable]).reshape(n_segments, nperseg, 2) + ground_truth = np.asarray(ground_truth_iq[:usable]).reshape( + n_segments, nperseg, 2 + ) + nmse = metrics.NMSE(prediction, ground_truth) + evm = metrics.EVM( + prediction, + ground_truth, + sample_rate=int(spec["input_signal_fs"]), + bw_main_ch=spec["bw_main_ch"], + n_sub_ch=spec["n_sub_ch"], + nperseg=nperseg, + ) + aclr_left, aclr_right = metrics.ACLR( + prediction, + fs=spec["input_signal_fs"], + nperseg=nperseg, + bw_main_ch=spec["bw_main_ch"], + n_sub_ch=spec["n_sub_ch"], + ) + return nmse, evm, aclr_left, aclr_right + + +def metric_dict(prediction_iq, ground_truth_iq, spec) -> dict[str, float]: + nmse, evm, aclr_left, aclr_right = compute_metrics( + prediction_iq, + ground_truth_iq, + spec, + ) + result = { + "ACLR_L": float(aclr_left), + "ACLR_R": float(aclr_right), + "ACLR_AVG": float((aclr_left + aclr_right) / 2.0), + "EVM": float(evm), + "NMSE": float(nmse), + } + if not all(math.isfinite(value) for value in result.values()): + raise RuntimeError(f"evaluation produced non-finite metrics: {result}") + return result + + +def display_metrics(split: str, values: dict[str, float]) -> None: + print(f"\n{split} results") + print(f"{'Metric':<12} {'Value':>12}") + print("-" * 26) + for name in ("ACLR_L", "ACLR_R", "ACLR_AVG", "EVM", "NMSE"): + print(f"{name:<12} {values[name]:>12.2f} dB") + + +def _portable_path(path: Path) -> str: + try: + return path.relative_to(REPO_ROOT).as_posix() + except ValueError: + return str(path) + + +# --------------------------------------------------------------------------- +# CLI and benchmark execution +# --------------------------------------------------------------------------- + +def build_arg_parser(): + parser = argparse.ArgumentParser( + description="Fit and evaluate MP/GMP PA or ILA-DPD benchmarks." + ) + parser.add_argument( + "--task", + choices=("pa_modeling", "dpd_ila"), + default="dpd_ila", + help="Benchmark task (default preserves the historical DPD behavior).", + ) + parser.add_argument( + "--dataset-name", + "--dataset_name", + dest="dataset_name", + required=True, + ) + parser.add_argument("--model", choices=("mp", "gmp"), default="mp") + + parser.add_argument("--K", type=int, default=5) + parser.add_argument("--Q", type=int, default=50) + parser.add_argument("--Ka", type=int, default=5) + parser.add_argument("--La", type=int, default=15) + parser.add_argument("--Kb", type=int, default=4) + parser.add_argument("--Lb", type=int, default=15) + parser.add_argument("--Mb", type=int, default=2) + parser.add_argument("--Kc", type=int, default=4) + parser.add_argument("--Lc", type=int, default=15) + parser.add_argument("--Mc", type=int, default=1) + + parser.add_argument( + "--solver-device", + default="cuda", + help="'cuda', 'cuda:INDEX', or 'cpu' (default: cuda).", + ) + parser.add_argument( + "--solver-dtype", + choices=("complex64", "complex128"), + default="complex64", + ) + parser.add_argument( + "--solver-mode", + choices=("gels", "truncated_svd"), + default="gels", + ) + parser.add_argument( + "--svd-rcond", + type=float, + help="Relative singular-value cutoff for truncated_svd.", + ) + parser.add_argument( + "--device", + type=int, + default=0, + help="CUDA index for PA evaluation; -1 selects CPU.", + ) + + parser.add_argument( + "--pa-backbone", + "--pa_backbone", + dest="pa_backbone", + default="tres_gru", + ) + parser.add_argument( + "--pa-hidden-size", + "--pa_hidden_size", + dest="pa_hidden_size", + type=int, + default=24, + ) + parser.add_argument( + "--pa-num-layers", + "--pa_num_layers", + dest="pa_num_layers", + type=int, + default=1, + ) + parser.add_argument("--pa-checkpoint", help="Explicit trained PA checkpoint.") + parser.add_argument("--frame-length", type=int, default=200) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--json-out", help="Write metrics and provenance as JSON.") + return parser + + +def run_benchmark(args: argparse.Namespace) -> dict: + spec_path = REPO_ROOT / "datasets" / args.dataset_name / "spec.json" + with open(spec_path, encoding="utf-8") as file: + spec = json.load(file) + nperseg = int(spec["nperseg"]) + configuration = basis_configuration(args) + _validate_positive_configuration(configuration) + n_complex_coefficients = coefficient_count(args.model, configuration) + n_real_parameters = 2 * n_complex_coefficients + solver_device = resolve_solver_device(args.solver_device, args.device) + solver_dtype = resolve_complex_dtype(args.solver_dtype) + + X_train, y_train, X_val, y_val, X_test, y_test = load_dataset( + dataset_name=args.dataset_name + ) + target_gain = float(set_target_gain(X_train, y_train)) + print( + f"\n=== {args.task}: {args.model.upper()} " + f"({n_complex_coefficients} complex / {n_real_parameters} real) ===" + ) + if args.solver_mode == "truncated_svd": + solver_description = ( + f"torch.linalg.svd truncated at rcond={args.svd_rcond:g}" + ) + else: + solver_description = "torch.linalg.lstsq gels" + print(f"Solver: {solver_description} on {solver_device} ({solver_dtype})") + + common_fit_arguments = { + "polynomial_model": args.model, + "configuration": configuration, + "segment_length": nperseg, + "device": solver_device, + "dtype": solver_dtype, + "solver_mode": args.solver_mode, + "svd_rcond": args.svd_rcond, + } + + if args.task == "pa_modeling": + coefficients, diagnostics = fit_complex_least_squares( + iq_to_complex(X_train), + iq_to_complex(y_train), + **common_fit_arguments, + ) + + validation_prediction = complex_to_iq( + apply_polynomial( + iq_to_complex(X_val), + coefficients, + polynomial_model=args.model, + configuration=configuration, + segment_length=nperseg, + ) + ) + validation_metrics = metric_dict(validation_prediction, y_val, spec) + del validation_prediction + + test_prediction = complex_to_iq( + apply_polynomial( + iq_to_complex(X_test), + coefficients, + polynomial_model=args.model, + configuration=configuration, + segment_length=nperseg, + ) + ) + test_metrics = metric_dict(test_prediction, y_test, spec) + method = "direct_least_squares" + pa_binding = {} + else: + coefficients, diagnostics = fit_complex_least_squares( + iq_to_complex(y_train) / target_gain, + iq_to_complex(X_train), + **common_fit_arguments, + ) + pa_net, pa_device, pa_path, pa_id, pa_parameters = load_pa_model(args) + + validation_dpd = complex_to_iq( + apply_polynomial( + iq_to_complex(X_val), + coefficients, + polynomial_model=args.model, + configuration=configuration, + segment_length=nperseg, + ) + ) + validation_prediction = run_pa_model( + validation_dpd, + pa_net, + pa_device, + nperseg, + ) + validation_target = (target_gain * X_val).astype(np.float32) + validation_metrics = metric_dict( + validation_prediction, + validation_target, + spec, + ) + del validation_dpd, validation_prediction + + test_dpd = complex_to_iq( + apply_polynomial( + iq_to_complex(X_test), + coefficients, + polynomial_model=args.model, + configuration=configuration, + segment_length=nperseg, + ) + ) + test_prediction = run_pa_model(test_dpd, pa_net, pa_device, nperseg) + test_target = (target_gain * X_test).astype(np.float32) + test_metrics = metric_dict(test_prediction, test_target, spec) + method = "indirect_learning_architecture" + pa_binding = { + "pa_backbone": args.pa_backbone, + "pa_hidden_size": int(args.pa_hidden_size), + "pa_num_layers": int(args.pa_num_layers), + "pa_model_id": pa_id, + "pa_parameters": int(pa_parameters), + "pa_evaluation_device": str(pa_device), + "pa_checkpoint": _portable_path(pa_path), + "pa_checkpoint_sha256": sha256_file(pa_path), + } + + display_metrics("Validation", validation_metrics) + display_metrics("Test", test_metrics) + dtype_name = str(solver_dtype) + device_name = str(solver_device) + result = { + "schema_version": 4, + "task": args.task, + "method": method, + "dataset": args.dataset_name, + "model": args.model, + "basis_configuration": configuration, + "complex_coefficients": int(n_complex_coefficients), + "real_parameters": int(n_real_parameters), + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "target_gain": target_gain, + "sample_rate_hz": float(spec["input_signal_fs"]), + "nperseg": nperseg, + "segment_boundary_policy": "zero delay state at every nperseg boundary", + "solver": diagnostics["solver_implementation"], + "solver_mode": diagnostics["solver_mode"], + "solver_driver": diagnostics["solver_driver"], + "device": device_name, + "dtype": dtype_name, + # Retain explicit aliases so downstream consumers cannot confuse the + # solver device with the separate neural-PA evaluation device. + "solver_device": device_name, + "solver_dtype": dtype_name, + "column_scaling": "l2", + "column_scale_min": diagnostics["column_norm_min"], + "column_scale_max": diagnostics["column_norm_max"], + "full_rank_assumption": args.solver_mode == "gels", + "least_squares_rank": diagnostics["least_squares_rank"], + "regularization": diagnostics["regularization"], + "svd_rcond": diagnostics["svd_rcond"], + "singular_value_max": diagnostics["singular_value_max"], + "singular_value_min": diagnostics["singular_value_min"], + "singular_value_cutoff": diagnostics["singular_value_cutoff"], + "retained_singular_value_min": diagnostics[ + "retained_singular_value_min" + ], + "condition_number": diagnostics["condition_number"], + "coefficient_l2_norm": diagnostics["coefficient_l2_norm"], + "training_relative_residual": diagnostics["training_relative_residual"], + "training_diagnostics": diagnostics, + "determinism": { + "stochastic_fitting": False, + "seed": int(args.seed), + "note": ( + "The design matrix and targets are deterministic; final low-order " + "floating-point bits can vary across CUDA library versions." + ), + }, + "validation": validation_metrics, + "test": test_metrics, + **pa_binding, + } + return result + + +def main(): + args = build_arg_parser().parse_args() + result = run_benchmark(args) + if args.json_out: + output_path = Path(args.json_out).expanduser() + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as file: + json.dump(result, file, indent=2, sort_keys=True) + file.write("\n") + print(f"\nRaw results written to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_volterra_qr.py b/benchmark/benchmark_volterra_qr.py deleted file mode 100644 index 0681789..0000000 --- a/benchmark/benchmark_volterra_qr.py +++ /dev/null @@ -1,302 +0,0 @@ -""" -Benchmark: Volterra-series-based DPD identified via QR decomposition. - -Supports two model structures: - - MP (Memory Polynomial): diagonal Volterra terms only - - GMP (Generalized Memory Polynomial): aligned + lagging + leading cross-terms - -Both identified in closed form via least squares (QR/SVD). - -Indirect Learning Architecture (ILA): - 1. Build basis Phi from normalized PA output z/G - 2. Solve x = Phi * w via QR (postdistorter) - 3. Copy w to predistorter: x_dpd = Phi(x) * w - 4. Pass through PA model: PA(x_dpd) vs G*x -""" - -import argparse -import json -import os -import numpy as np -import torch - -# Project imports -from modules.data_collector import load_dataset -from utils.util import set_target_gain -from utils import metrics -import models as model - - -# --------------------------------------------------------------------------- -# Memory Polynomial basis -# --------------------------------------------------------------------------- - -def build_mp_basis(x_complex, K, Q): - """ - Build Memory Polynomial basis matrix. - - Basis functions: phi_{k,q}(n) = x(n-q) * |x(n-q)|^k - for k = 0..K-1, q = 0..Q-1 - - Parameters - ---------- - x_complex : ndarray, shape (N,), complex - K : int - number of nonlinearity orders - Q : int - memory depth - - Returns - ------- - Phi : ndarray, shape (N, K*Q), complex - """ - N = len(x_complex) - Phi = np.zeros((N, K * Q), dtype=np.complex128) - col = 0 - for k in range(K): - for q in range(Q): - if q == 0: - xd = x_complex - else: - xd = np.zeros(N, dtype=np.complex128) - xd[q:] = x_complex[:N - q] - Phi[:, col] = xd * np.abs(xd) ** k - col += 1 - return Phi - - -# --------------------------------------------------------------------------- -# GMP basis (Morgan et al., IEEE TSP 2006) -# --------------------------------------------------------------------------- - -def _delay(x, d): - """Delay complex signal x by d samples (zero-pad).""" - N = len(x) - out = np.zeros(N, dtype=np.complex128) - if d >= 0 and d < N: - out[d:] = x[:N - d] - elif d < 0 and -d < N: - out[:N + d] = x[-d:] - return out - - -def build_gmp_basis(x_complex, Ka, La, Kb, Lb, Mb, Kc, Lc, Mc): - """ - Build Generalized Memory Polynomial basis matrix. - - Three components: - Aligned: x(n-q) * |x(n-q)|^k k=0..Ka-1, q=0..La-1 - Lagging: x(n-q) * |x(n-q-l)|^k k=1..Kb, q=0..Lb-1, l=1..Mb - Leading: x(n-q) * |x(n-q+l)|^k k=1..Kc, q=0..Lc-1, l=1..Mc - - Returns Phi of shape (N, Ka*La + Kb*Lb*Mb + Kc*Lc*Mc), complex. - """ - N = len(x_complex) - n_cols = Ka * La + Kb * Lb * Mb + Kc * Lc * Mc - Phi = np.zeros((N, n_cols), dtype=np.complex128) - col = 0 - - # Aligned terms - for k in range(Ka): - for q in range(La): - xq = _delay(x_complex, q) - Phi[:, col] = xq * np.abs(xq) ** k - col += 1 - - # Lagging cross-terms - for k in range(1, Kb + 1): - for q in range(Lb): - for l in range(1, Mb + 1): - xq = _delay(x_complex, q) - xql = _delay(x_complex, q + l) - Phi[:, col] = xq * np.abs(xql) ** k - col += 1 - - # Leading cross-terms - for k in range(1, Kc + 1): - for q in range(Lc): - for l in range(1, Mc + 1): - xq = _delay(x_complex, q) - xql = _delay(x_complex, q - l) - Phi[:, col] = xq * np.abs(xql) ** k - col += 1 - - return Phi - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -def iq_to_complex(iq): - return iq[:, 0] + 1j * iq[:, 1] - - -def complex_to_iq(c): - return np.stack([c.real, c.imag], axis=-1).astype(np.float32) - - -def run_pa_model(iq_signal, pa_net, device, nperseg): - """Pass an IQ signal through the trained PA model segment-by-segment.""" - pa_net.eval() - N = len(iq_signal) - segments = [] - for i in range(0, N, nperseg): - seg = iq_signal[i:i + nperseg] - if len(seg) < nperseg: - pad = np.zeros((nperseg - len(seg), 2), dtype=np.float32) - seg = np.vstack([seg, pad]) - segments.append(seg) - segments = np.array(segments) # (n_seg, nperseg, 2) - - with torch.no_grad(): - out = pa_net(torch.tensor(segments, dtype=torch.float32).to(device)) - out = out.cpu().numpy() - - return out.reshape(-1, 2)[:N] - - -def compute_metrics(prediction_iq, ground_truth_iq, spec): - """Compute NMSE, EVM, ACLR on segmented data (same as OpenDPD training).""" - nperseg = int(spec['nperseg']) - fs = spec['input_signal_fs'] - bw = spec['bw_main_ch'] - nsc = spec['n_sub_ch'] - - # Segment prediction and ground truth - N = len(prediction_iq) - n_seg = N // nperseg - pred_seg = prediction_iq[:n_seg * nperseg].reshape(n_seg, nperseg, 2) - gt_seg = ground_truth_iq[:n_seg * nperseg].reshape(n_seg, nperseg, 2) - - nmse = metrics.NMSE(pred_seg, gt_seg) - evm = metrics.EVM(pred_seg, gt_seg, sample_rate=int(fs), - bw_main_ch=bw, n_sub_ch=nsc, nperseg=nperseg) - aclr_l, aclr_r = metrics.ACLR(pred_seg, fs=fs, nperseg=nperseg, - bw_main_ch=bw, n_sub_ch=nsc) - return nmse, evm, aclr_l, aclr_r - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--dataset_name', type=str, required=True) - parser.add_argument('--pa_hidden_size', type=int, default=24) - parser.add_argument('--model', type=str, default='mp', choices=['mp', 'gmp'], - help='Model type: mp (Memory Polynomial) or gmp (Generalized MP)') - # MP parameters - parser.add_argument('--K', type=int, default=5, - help='Nonlinearity order (number of |x|^k terms)') - parser.add_argument('--Q', type=int, default=50, - help='Memory depth (number of delay taps) for MP') - # GMP parameters - parser.add_argument('--Ka', type=int, default=5, help='GMP aligned nonlinearity orders') - parser.add_argument('--La', type=int, default=15, help='GMP aligned memory depth') - parser.add_argument('--Kb', type=int, default=4, help='GMP lagging nonlinearity orders') - parser.add_argument('--Lb', type=int, default=15, help='GMP lagging memory depth') - parser.add_argument('--Mb', type=int, default=2, help='GMP lagging cross-term depth') - parser.add_argument('--Kc', type=int, default=4, help='GMP leading nonlinearity orders') - parser.add_argument('--Lc', type=int, default=15, help='GMP leading memory depth') - parser.add_argument('--Mc', type=int, default=1, help='GMP leading cross-term depth') - parser.add_argument('--device', type=int, default=0) - args = parser.parse_args() - - # Load spec - spec_path = os.path.join('datasets', args.dataset_name, 'spec.json') - with open(spec_path) as f: - spec = json.load(f) - nperseg = int(spec['nperseg']) - - # Load dataset - X_train, y_train, X_val, y_val, X_test, y_test = load_dataset( - dataset_name=args.dataset_name) - - # Target gain (same as OpenDPD) - target_gain = set_target_gain(X_train, y_train) - - # Load PA model - device = torch.device(f'cuda:{args.device}') - pa_net = model.CoreModel(input_size=2, hidden_size=args.pa_hidden_size, - num_layers=1, backbone_type='gru') - from utils.util import count_net_params - n_pa = count_net_params(pa_net) - pa_path = (f'save/{args.dataset_name}/train_pa/' - f'PA_S_0_M_GRU_H_{args.pa_hidden_size}_F_200_P_{n_pa}.pt') - pa_net.load_state_dict(torch.load(pa_path, map_location='cpu')) - pa_net = pa_net.to(device) - - # ----------------------------------------------------------------------- - # Step 1: Identify DPD via ILA + QR - # ----------------------------------------------------------------------- - K, Q = args.K, args.Q - n_complex_coeffs = K * Q - n_real_params = 2 * n_complex_coeffs - print(f"\n=== Volterra MP DPD via QR (K={K}, Q={Q}) ===") - print(f"Complex coefficients : {n_complex_coeffs}") - print(f"Real parameters : {n_real_params}") - - # Convert to complex - x_train_c = iq_to_complex(X_train) - z_train_c = iq_to_complex(y_train) # measured PA output - - # Normalise PA output by target gain - z_norm = z_train_c / target_gain - - # Build basis from normalised PA output (postdistorter input) - Phi_train = build_mp_basis(z_norm, K, Q) - - # Solve x_train = Phi_train * w via least squares (uses QR/SVD internally) - w, residuals, rank, sv = np.linalg.lstsq(Phi_train, x_train_c, rcond=None) - print(f"Least-squares rank : {rank} / {n_complex_coeffs}") - - # ----------------------------------------------------------------------- - # Step 2: Apply DPD to test data - # ----------------------------------------------------------------------- - x_test_c = iq_to_complex(X_test) - - # Build basis from test input (predistorter) - Phi_test = build_mp_basis(x_test_c, K, Q) - x_dpd_c = Phi_test @ w # predistorted signal - x_dpd_iq = complex_to_iq(x_dpd_c) - - # ----------------------------------------------------------------------- - # Step 3: Pass through PA model - # ----------------------------------------------------------------------- - pa_output_iq = run_pa_model(x_dpd_iq, pa_net, device, nperseg) - - # Ground truth for DPD: target_gain * x_test - gt_iq = (target_gain * X_test).astype(np.float32) - - # ----------------------------------------------------------------------- - # Step 4: Metrics - # ----------------------------------------------------------------------- - nmse, evm, aclr_l, aclr_r = compute_metrics(pa_output_iq, gt_iq, spec) - aclr_avg = (aclr_l + aclr_r) / 2 - - print(f"\nTest Results:") - print(f"{'Metric':<12} {'Value':>12}") - print("-" * 26) - print(f"{'ACLR_L':<12} {aclr_l:>12.2f} dB") - print(f"{'ACLR_R':<12} {aclr_r:>12.2f} dB") - print(f"{'ACLR_AVG':<12} {aclr_avg:>12.2f} dB") - print(f"{'EVM':<12} {evm:>12.2f} dB") - print(f"{'NMSE':<12} {nmse:>12.2f} dB") - - # Also run on val set - x_val_c = iq_to_complex(X_val) - Phi_val = build_mp_basis(x_val_c, K, Q) - x_dpd_val_c = Phi_val @ w - x_dpd_val_iq = complex_to_iq(x_dpd_val_c) - pa_val_iq = run_pa_model(x_dpd_val_iq, pa_net, device, nperseg) - gt_val_iq = (target_gain * X_val).astype(np.float32) - v_nmse, v_evm, v_aclr_l, v_aclr_r = compute_metrics(pa_val_iq, gt_val_iq, spec) - v_aclr_avg = (v_aclr_l + v_aclr_r) / 2 - print(f"\nVal Results:") - print(f"{'ACLR_AVG':<12} {v_aclr_avg:>12.2f} dB") - print(f"{'EVM':<12} {v_evm:>12.2f} dB") - print(f"{'NMSE':<12} {v_nmse:>12.2f} dB") - - -if __name__ == '__main__': - main() diff --git a/benchmark/collect_benchmark_report.py b/benchmark/collect_benchmark_report.py new file mode 100644 index 0000000..9e78f51 --- /dev/null +++ b/benchmark/collect_benchmark_report.py @@ -0,0 +1,3011 @@ +"""Collect one benchmark-report reproduction run into JSON and Markdown. + +The training entry point writes checkpoints and CSV logs to architecture-based +paths under ``save/`` and ``log/``. The benchmark runner archives those files +into a run-specific directory, then calls this module to validate the expected +recipes and assemble a portable evidence bundle. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +from datetime import datetime, timezone +import gzip +import hashlib +import importlib.metadata +import io +import json +import math +import os +from pathlib import Path +import platform +import shlex +import subprocess +import tarfile +import tempfile +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parent.parent +RUN_SCHEMA_VERSION = 6 +POLYNOMIAL_SCHEMA_VERSION = 4 +GMP_PA_SVD_RCOND = 1e-4 +PA_MODEL_ORDER = ("mp", "gmp", "gru", "tres_gru", "tres_deltagru") +DPD_MODEL_ORDER = ("mp", "gmp", "gru", "tres_gru") + +DATASETS = { + "APA_200MHz": { + "slug": "apa", + "sample_rate_hz": 983_040_000.0, + "nperseg": 19_662, + }, + "DPA_160MHz": { + "slug": "dpa", + "sample_rate_hz": 640_000_000.0, + "nperseg": 16_384, + }, +} + +NEURAL_PA_MODELS = { + "gru": { + "display_name": "GRU-H28", + "backbone": "gru", + "hidden_size": 28, + "parameters": 2746, + "model_id": "PA_S_0_M_GRU_H_28_F_200_P_2746", + }, + "tres_gru": { + "display_name": "TRes-GRU-H27", + "backbone": "tres_gru", + "hidden_size": 27, + "parameters": 2751, + "model_id": "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751", + }, + "tres_deltagru": { + "display_name": "TRes-DeltaGRU-H27 (THX=THH=0)", + "backbone": "tres_deltagru", + "hidden_size": 27, + "parameters": 2751, + "model_id": "PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751", + "delta_thresholds": { + "input": 0.0, + "hidden": 0.0, + }, + }, +} + +REFERENCE_PA_KEY = "tres_gru" +REFERENCE_PA_PARENT = "PA_S_0_M_TRES_GRU_H_27_F_200" +ALTERNATE_PA_KEY = "tres_deltagru" +ALTERNATE_PA_PARENT = "PA_S_0_M_TRES_DELTAGRU_H_27_F_200" + +NEURAL_DPD_MODELS = { + "gru": { + "display_name": "GRU-H16", + "backbone": "gru", + "hidden_size": 16, + "parameters": 994, + "model_id": "DPD_S_0_M_GRU_H_16_F_200_P_994", + }, + "tres_gru": { + "display_name": "TRes-GRU-H15", + "backbone": "tres_gru", + "hidden_size": 15, + "parameters": 999, + "model_id": "DPD_S_0_M_TRES_GRU_H_15_F_200_P_999", + }, +} + +TRES_DELTAGRU_DPD_MODEL = { + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "backbone": "tres_deltagru", + "hidden_size": 15, + "parameters": 999, + "model_id": ( + "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_" + "THX_0.000_THH_0.000" + ), + "delta_thresholds": { + "input": 0.0, + "hidden": 0.0, + }, +} + +NEURAL_DPD_EXPERIMENTS = { + **{ + f"{model_key}_via_{REFERENCE_PA_KEY}_pa": { + "model_key": model_key, + "pa_model_key": REFERENCE_PA_KEY, + "pa_parent": REFERENCE_PA_PARENT, + } + for model_key in NEURAL_DPD_MODELS + }, + "tres_deltagru_via_tres_gru_pa": { + "model_key": "tres_deltagru", + "pa_model_key": REFERENCE_PA_KEY, + "pa_parent": REFERENCE_PA_PARENT, + }, + "tres_deltagru_via_tres_deltagru_pa": { + "model_key": "tres_deltagru", + "pa_model_key": ALTERNATE_PA_KEY, + "pa_parent": ALTERNATE_PA_PARENT, + }, +} + +POLYNOMIAL_CONFIGURATIONS = { + "pa_modeling": { + "mp": { + "basis_configuration": {"K": 9, "Q": 150}, + "complex_coefficients": 1350, + "real_parameters": 2700, + }, + "gmp": { + "basis_configuration": { + "Ka": 5, + "La": 30, + "Kb": 4, + "Lb": 30, + "Mb": 5, + "Kc": 4, + "Lc": 30, + "Mc": 5, + }, + "complex_coefficients": 1350, + "real_parameters": 2700, + }, + }, + "dpd_ila": { + "mp": { + "basis_configuration": {"K": 5, "Q": 100}, + "complex_coefficients": 500, + "real_parameters": 1000, + }, + "gmp": { + "basis_configuration": { + "Ka": 5, + "La": 20, + "Kb": 4, + "Lb": 20, + "Mb": 3, + "Kc": 4, + "Lc": 20, + "Mc": 2, + }, + "complex_coefficients": 500, + "real_parameters": 1000, + }, + }, +} + +NEURAL_RECIPE = { + "epochs": 300, + "batch_size": 64, + "initial_learning_rate": 5e-3, + "lr_schedule": "ReduceLROnPlateau", + "minimum_learning_rate": 5e-5, + "decay_factor": 0.5, + "patience": 5, + "frame_length": 200, + "frame_stride": 1, + "optimizer": "AdamW", + "optimizer_weight_decay": 0.01, + "optimizer_betas": [0.9, 0.999], + "optimizer_eps": 1e-8, + "loss": "MSE", + "scheduler_threshold": 1e-4, + "scheduler_threshold_mode": "rel", + "scheduler_cooldown": 0, + "scheduler_eps": 1e-8, + "seed": 0, +} + +SOURCE_PATTERNS = ( + "*.py", + "backbones/**/*.py", + "benchmark/**/*.py", + "modules/**/*.py", + "opendpd/**/*.py", + "quant/**/*.py", + "steps/**/*.py", + "utils/**/*.py", +) +SOURCE_DOCUMENTS = ( + "benchmark/benchmark_report.md", + "benchmark/reproduce_benchmark_report.sh", +) + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def write_source_snapshot(paths: list[str], destination: Path) -> None: + """Write an atomic, deterministic archive of the exact benchmark sources.""" + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=destination.parent, + prefix=f".{destination.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + with temporary_path.open("wb") as raw_handle: + with gzip.GzipFile( + filename="", + mode="wb", + fileobj=raw_handle, + mtime=0, + ) as gzip_handle: + with tarfile.open( + fileobj=gzip_handle, + mode="w", + format=tarfile.GNU_FORMAT, + ) as archive: + for relative_path in sorted(paths): + source = REPO_ROOT / relative_path + if not source.is_file() or source.is_symlink(): + raise ValueError( + "Source snapshot entries must be regular files: " + f"{relative_path}" + ) + content = source.read_bytes() + information = tarfile.TarInfo(relative_path) + information.size = len(content) + information.mode = source.stat().st_mode & 0o777 + information.mtime = 0 + information.uid = 0 + information.gid = 0 + information.uname = "" + information.gname = "" + archive.addfile(information, io.BytesIO(content)) + os.replace(temporary_path, destination) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() + + +def validate_source_snapshot( + path: Path, + expected_hashes: dict[str, str], +) -> dict[str, Any]: + """Verify every archive member against the pre-run source hash map.""" + if not path.is_file(): + raise FileNotFoundError(f"Source snapshot is missing: {path}") + observed: dict[str, str] = {} + with tarfile.open(path, mode="r:gz") as archive: + for member in archive.getmembers(): + if not member.isfile(): + raise ValueError( + f"Source snapshot contains a non-file member: {member.name}" + ) + if member.name in observed: + raise ValueError( + f"Source snapshot contains a duplicate member: {member.name}" + ) + handle = archive.extractfile(member) + if handle is None: + raise ValueError( + f"Cannot read source snapshot member: {member.name}" + ) + observed[member.name] = hashlib.sha256(handle.read()).hexdigest() + require_equal(observed, expected_hashes, "source snapshot members") + return { + "path": path.name, + "sha256": sha256_file(path), + "file_count": len(observed), + } + + +def read_csv(path: Path) -> list[dict[str, str]]: + with path.open(newline="") as handle: + rows = list(csv.DictReader(handle)) + if not rows: + raise ValueError(f"CSV contains no data rows: {path}") + return rows + + +def as_int(row: dict[str, str], key: str) -> int: + value = as_float(row, key) + if not value.is_integer(): + raise ValueError(f"{key} is not an integer: {value}") + return int(value) + + +def as_float(row: dict[str, str], key: str) -> float: + value = float(row[key]) + if not math.isfinite(value): + raise ValueError(f"{key} is not finite: {value}") + return value + + +def require_equal(actual: Any, expected: Any, description: str) -> None: + if actual != expected: + raise ValueError( + f"{description} does not match the benchmark recipe: " + f"expected {expected!r}, got {actual!r}" + ) + + +def require_close( + actual: float, + expected: float, + description: str, + *, + tolerance: float = 1e-12, +) -> None: + if not math.isclose(actual, expected, rel_tol=0.0, abs_tol=tolerance): + raise ValueError( + f"{description} does not match the benchmark recipe: " + f"expected {expected}, got {actual}" + ) + + +def metric_block(row: dict[str, str]) -> dict[str, dict[str, float]]: + def split(prefix: str) -> dict[str, float]: + metrics = { + "loss": as_float(row, f"{prefix}_LOSS"), + "nmse_db": as_float(row, f"{prefix}_NMSE"), + "evm_db": as_float(row, f"{prefix}_EVM"), + "aclr_left_db": as_float(row, f"{prefix}_ACLR_L"), + "aclr_right_db": as_float(row, f"{prefix}_ACLR_R"), + "aclr_avg_db": as_float(row, f"{prefix}_ACLR_AVG"), + } + require_close( + metrics["aclr_avg_db"], + (metrics["aclr_left_db"] + metrics["aclr_right_db"]) / 2.0, + f"{prefix} logged ACLR average", + tolerance=1e-8, + ) + return metrics + + return {"validation": split("VAL"), "test": split("TEST")} + + +def source_files() -> tuple[str, ...]: + paths = { + path.relative_to(REPO_ROOT).as_posix() + for pattern in SOURCE_PATTERNS + for path in REPO_ROOT.glob(pattern) + if path.is_file() + } + paths.update(SOURCE_DOCUMENTS) + missing = [path for path in paths if not (REPO_ROOT / path).is_file()] + if missing: + raise FileNotFoundError( + f"Expected benchmark source files are missing: {sorted(missing)}" + ) + return tuple(sorted(paths)) + + +def validate_checkpoint(path: Path, expected_parameters: int) -> None: + import torch + + try: + state_dict = torch.load(path, map_location="cpu", weights_only=True) + except Exception as exc: + raise ValueError(f"Checkpoint cannot be loaded safely: {path}") from exc + if not isinstance(state_dict, dict) or not state_dict: + raise ValueError(f"Checkpoint is not a non-empty state dictionary: {path}") + if not all( + isinstance(key, str) and isinstance(value, torch.Tensor) + for key, value in state_dict.items() + ): + raise ValueError(f"Checkpoint contains non-tensor state entries: {path}") + parameter_count = sum(value.numel() for value in state_dict.values()) + require_equal( + parameter_count, + expected_parameters, + f"checkpoint parameter count for {path}", + ) + if not all(torch.isfinite(value).all().item() for value in state_dict.values()): + raise ValueError(f"Checkpoint contains a non-finite tensor: {path}") + + +def archived_artifact(output_dir: Path, source: Path) -> dict[str, Any]: + absolute_source = REPO_ROOT / source + archived_path = output_dir / "artifacts" / source + if not absolute_source.is_file(): + raise FileNotFoundError(f"Expected benchmark artifact is missing: {source}") + if not archived_path.is_file(): + raise FileNotFoundError( + f"Benchmark artifact was not archived into the run directory: {source}" + ) + source_hash = sha256_file(absolute_source) + archived_hash = sha256_file(archived_path) + require_equal(archived_hash, source_hash, f"archived hash for {source}") + return { + "source_path": source.as_posix(), + "archived_path": archived_path.relative_to(output_dir).as_posix(), + "sha256": source_hash, + "size_bytes": absolute_source.stat().st_size, + } + + +def read_pa_bindings( + output_dir: Path, +) -> dict[tuple[str, str, str], dict[str, str]]: + path = output_dir / "pa_checkpoint_bindings.tsv" + with path.open(newline="") as handle: + rows = list(csv.DictReader(handle, delimiter="\t")) + require_equal( + len(rows), + len(DATASETS) * len(NEURAL_DPD_EXPERIMENTS), + "neural DPD PA checkpoint binding count", + ) + + bindings: dict[tuple[str, str, str], dict[str, str]] = {} + for row in rows: + key = ( + row["dataset"], + row["dpd_model_key"], + row["pa_model_key"], + ) + if key in bindings: + raise ValueError(f"Duplicate PA checkpoint binding: {key}") + if row["dataset"] not in DATASETS: + raise ValueError(f"Unknown dataset in PA checkpoint binding: {key}") + if ( + row["dpd_model_key"] not in NEURAL_DPD_MODELS + and row["dpd_model_key"] != "tres_deltagru" + ): + raise ValueError(f"Unknown model in PA checkpoint binding: {key}") + if row["pa_model_key"] not in {REFERENCE_PA_KEY, ALTERNATE_PA_KEY}: + raise ValueError(f"Unknown PA surrogate in checkpoint binding: {key}") + if len(row["sha256"]) != 64: + raise ValueError(f"Invalid PA checkpoint hash for {key}") + bindings[key] = row + + expected_keys = { + ( + dataset, + experiment["model_key"], + experiment["pa_model_key"], + ) + for dataset in DATASETS + for experiment in NEURAL_DPD_EXPERIMENTS.values() + } + require_equal(set(bindings), expected_keys, "neural DPD PA checkpoint bindings") + return bindings + + +def read_recorded_commands(path: Path) -> dict[str, list[str]]: + commands: dict[str, list[str]] = {} + current_label: str | None = None + for raw_line in path.read_text().splitlines(): + line = raw_line.strip() + if line.startswith("# "): + current_label = line[2:] + if current_label in commands: + raise ValueError(f"Duplicate command label in {path}: {current_label}") + continue + if not line: + continue + if current_label is None: + raise ValueError(f"Unlabelled command in {path}: {line}") + commands[current_label] = shlex.split(line) + current_label = None + if current_label is not None: + raise ValueError(f"Command label has no command in {path}: {current_label}") + return commands + + +def cli_options(command: list[str], *, label: str) -> dict[str, str]: + try: + main_index = command.index("main.py") + except ValueError as exc: + raise ValueError(f"Recorded training command lacks main.py: {label}") from exc + tokens = command[main_index + 1 :] + if len(tokens) % 2: + raise ValueError(f"Recorded training command has an unpaired option: {label}") + options: dict[str, str] = {} + for option, value in zip(tokens[::2], tokens[1::2]): + if not option.startswith("--"): + raise ValueError( + f"Unexpected positional token in recorded training command " + f"{label}: {option}" + ) + key = option[2:] + if key in options: + raise ValueError(f"Duplicate --{key} in recorded command: {label}") + options[key] = value + return options + + +def validate_neural_commands(path: Path, device: int | None = None) -> None: + commands = read_recorded_commands(path) + training_commands = { + label: cli_options(command, label=label) + for label, command in commands.items() + if "main.py" in command + } + require_equal( + len(training_commands), + len(DATASETS) * (len(NEURAL_PA_MODELS) + len(NEURAL_DPD_EXPERIMENTS)), + "recorded neural training command count", + ) + + expected_matrix = { + *( + ( + "train_pa", + dataset, + model["backbone"], + model["hidden_size"], + None, + None, + ) + for dataset in DATASETS + for model in NEURAL_PA_MODELS.values() + ), + *( + ( + "train_dpd", + dataset, + ( + NEURAL_DPD_MODELS[experiment["model_key"]] + if experiment["model_key"] in NEURAL_DPD_MODELS + else TRES_DELTAGRU_DPD_MODEL + )["backbone"], + ( + NEURAL_DPD_MODELS[experiment["model_key"]] + if experiment["model_key"] in NEURAL_DPD_MODELS + else TRES_DELTAGRU_DPD_MODEL + )["hidden_size"], + NEURAL_PA_MODELS[experiment["pa_model_key"]]["backbone"], + NEURAL_PA_MODELS[experiment["pa_model_key"]]["hidden_size"], + ) + for dataset in DATASETS + for experiment in NEURAL_DPD_EXPERIMENTS.values() + ), + } + actual_matrix: set[tuple[str, str, str, int, str | None, int | None]] = set() + for label, options in training_commands.items(): + step = options.get("step") + if step not in {"train_pa", "train_dpd"}: + raise ValueError(f"Unexpected training step in {label}: {step!r}") + backbone_key = "PA_backbone" if step == "train_pa" else "DPD_backbone" + hidden_key = "PA_hidden_size" if step == "train_pa" else "DPD_hidden_size" + try: + matrix_key = ( + step, + options["dataset_name"], + options[backbone_key], + int(options[hidden_key]), + options.get("PA_backbone") if step == "train_dpd" else None, + ( + int(options["PA_hidden_size"]) + if step == "train_dpd" + else None + ), + ) + except (KeyError, ValueError) as exc: + raise ValueError(f"Incomplete model recipe in command {label}") from exc + if matrix_key in actual_matrix: + raise ValueError(f"Duplicate neural benchmark command: {matrix_key}") + actual_matrix.add(matrix_key) + + integer_expectations = { + "n_epochs": NEURAL_RECIPE["epochs"], + "batch_size": NEURAL_RECIPE["batch_size"], + "batch_size_eval": NEURAL_RECIPE["batch_size"], + "lr_schedule": 1, + "frame_length": NEURAL_RECIPE["frame_length"], + "frame_stride": NEURAL_RECIPE["frame_stride"], + "PA_num_layers": 1, + "DPD_num_layers": 1, + "log_precision": 8, + "eval_val": 1, + "eval_test": 1, + "seed": NEURAL_RECIPE["seed"], + } + if device is not None: + integer_expectations["devices"] = device + for option, expected in integer_expectations.items(): + try: + actual = int(options[option]) + except (KeyError, ValueError) as exc: + raise ValueError( + f"Missing or invalid --{option} in command {label}" + ) from exc + require_equal(actual, expected, f"--{option} in command {label}") + + float_expectations = { + "lr": NEURAL_RECIPE["initial_learning_rate"], + "lr_end": NEURAL_RECIPE["minimum_learning_rate"], + "decay_factor": NEURAL_RECIPE["decay_factor"], + "patience": NEURAL_RECIPE["patience"], + "grad_clip_val": 200.0, + "thx": 0.0, + "thh": 0.0, + } + for option, expected in float_expectations.items(): + try: + actual = float(options[option]) + except (KeyError, ValueError) as exc: + raise ValueError( + f"Missing or invalid --{option} in command {label}" + ) from exc + require_close(actual, expected, f"--{option} in command {label}") + + require_equal(options.get("opt_type"), "adamw", f"optimizer in {label}") + require_equal(options.get("loss_type"), "l2", f"loss in {label}") + require_equal(options.get("accelerator"), "cuda", f"accelerator in {label}") + require_equal( + options.get("re_level"), + "soft", + f"reproducibility level in {label}", + ) + require_equal( + actual_matrix, + expected_matrix, + "recorded neural model matrix", + ) + + +def module_cli_options( + command: list[str], + *, + label: str, + module: str, +) -> dict[str, str]: + try: + module_index = command.index("-m") + except ValueError as exc: + raise ValueError(f"Recorded command lacks -m {module}: {label}") from exc + if ( + module_index + 1 >= len(command) + or command[module_index + 1] != module + ): + raise ValueError(f"Recorded command uses the wrong module: {label}") + tokens = command[module_index + 2 :] + if len(tokens) % 2: + raise ValueError(f"Recorded module command has an unpaired option: {label}") + options: dict[str, str] = {} + for option, value in zip(tokens[::2], tokens[1::2]): + if not option.startswith("--"): + raise ValueError( + f"Unexpected positional token in recorded command {label}: " + f"{option}" + ) + key = option[2:] + if key in options: + raise ValueError(f"Duplicate --{key} in recorded command: {label}") + options[key] = value + return options + + +def validate_runner_commands( + path: Path, + *, + output_dir: Path, + device: int, +) -> None: + """Validate the complete recorded neural and polynomial experiment matrix.""" + validate_neural_commands(path, device=device) + commands = read_recorded_commands(path) + expected_labels = {"snapshot_context", "collect_report"} + expected_labels.update( + f"train_pa_{configuration['slug']}_{model}" + for configuration in DATASETS.values() + for model in NEURAL_PA_MODELS + ) + expected_labels.update( + f"train_dpd_{configuration['slug']}_{experiment_key}" + for configuration in DATASETS.values() + for experiment_key in NEURAL_DPD_EXPERIMENTS + ) + expected_labels.update( + f"pa_model_{configuration['slug']}_{model}" + for configuration in DATASETS.values() + for model in ("mp", "gmp") + ) + expected_labels.update( + f"dpd_ila_{configuration['slug']}_{model}" + for configuration in DATASETS.values() + for model in ("mp", "gmp") + ) + require_equal(set(commands), expected_labels, "recorded benchmark job matrix") + + for label in ("snapshot_context", "collect_report"): + command = commands[label] + require_equal( + command.count("benchmark.collect_benchmark_report"), + 1, + f"collector module in {label}", + ) + require_equal( + command[command.index("--output-dir") + 1], + str(output_dir), + f"collector output directory in {label}", + ) + require_equal( + int(command[command.index("--device") + 1]), + device, + f"collector device in {label}", + ) + require_equal( + "--snapshot-context" in commands["snapshot_context"], + True, + "pre-run context command", + ) + require_equal( + "--snapshot-context" in commands["collect_report"], + False, + "final collector command", + ) + + expected_solver_device = f"cuda:{device}" + for dataset, dataset_configuration in DATASETS.items(): + slug = dataset_configuration["slug"] + for task, label_prefix, filename_task in ( + ("pa_modeling", "pa_model", "pa"), + ("dpd_ila", "dpd_ila", "dpd"), + ): + for model_name in ("mp", "gmp"): + label = f"{label_prefix}_{slug}_{model_name}" + options = module_cli_options( + commands[label], + label=label, + module="benchmark.benchmark_volterra", + ) + expectations = { + "task": task, + "dataset-name": dataset, + "model": model_name, + "solver-device": expected_solver_device, + "solver-dtype": "complex64", + "device": str(device), + "json-out": str( + output_dir + / ( + f"benchmark_report_{slug}_{filename_task}_" + f"{model_name}.json" + ) + ), + } + if task == "pa_modeling" and model_name == "gmp": + expectations.update( + { + "solver-mode": "truncated_svd", + "svd-rcond": "1e-4", + } + ) + else: + expectations["solver-mode"] = "gels" + require_equal( + "svd-rcond" in options, + False, + f"unexpected --svd-rcond in command {label}", + ) + if task == "dpd_ila": + expectations.update( + { + "pa-backbone": "tres_gru", + "pa-hidden-size": "27", + "pa-num-layers": "1", + "pa-checkpoint": ( + f"save/{dataset}/train_pa/" + "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt" + ), + } + ) + for key, value in expectations.items(): + require_equal( + options.get(key), + value, + f"--{key} in command {label}", + ) + + configuration = POLYNOMIAL_CONFIGURATIONS[task][model_name][ + "basis_configuration" + ] + for key, value in configuration.items(): + require_equal( + int(options[key]), + value, + f"--{key} in command {label}", + ) + + +def expected_runner_recipe() -> dict[str, Any]: + return { + "schema_version": RUN_SCHEMA_VERSION, + "datasets": list(DATASETS), + "parameter_count_convention": ( + "Neural parameters are real trainable scalars; each complex " + "polynomial coefficient counts as two real degrees of freedom." + ), + "neural": { + "optimizer": "AdamW", + "optimizer_weight_decay": NEURAL_RECIPE["optimizer_weight_decay"], + "optimizer_betas": NEURAL_RECIPE["optimizer_betas"], + "optimizer_eps": NEURAL_RECIPE["optimizer_eps"], + "loss": "MSE", + "epochs": NEURAL_RECIPE["epochs"], + "batch_size": NEURAL_RECIPE["batch_size"], + "evaluation_batch_size": NEURAL_RECIPE["batch_size"], + "initial_learning_rate": NEURAL_RECIPE["initial_learning_rate"], + "scheduler": "ReduceLROnPlateau", + "scheduler_factor": NEURAL_RECIPE["decay_factor"], + "scheduler_patience": NEURAL_RECIPE["patience"], + "scheduler_threshold": NEURAL_RECIPE["scheduler_threshold"], + "scheduler_threshold_mode": NEURAL_RECIPE[ + "scheduler_threshold_mode" + ], + "scheduler_cooldown": NEURAL_RECIPE["scheduler_cooldown"], + "scheduler_eps": NEURAL_RECIPE["scheduler_eps"], + "minimum_learning_rate": NEURAL_RECIPE[ + "minimum_learning_rate" + ], + "frame_length": NEURAL_RECIPE["frame_length"], + "frame_stride": NEURAL_RECIPE["frame_stride"], + "gradient_clip": 200, + "seed": NEURAL_RECIPE["seed"], + "reproducibility": "soft", + "pa_scheduler_and_selection_metric": "validation NMSE", + "dpd_scheduler_and_selection_metric": "validation ACLR_AVG", + }, + "pa_models": { + "mp": {"K": 9, "Q": 150, "real_parameters": 2700}, + "gmp": { + **POLYNOMIAL_CONFIGURATIONS["pa_modeling"]["gmp"][ + "basis_configuration" + ], + "real_parameters": 2700, + }, + "gru": {"hidden_size": 28, "real_parameters": 2746}, + "tres_gru": {"hidden_size": 27, "real_parameters": 2751}, + "tres_deltagru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": { + "input": 0.0, + "hidden": 0.0, + }, + }, + }, + "dpd_models": { + "mp": { + "K": 5, + "Q": 100, + "real_parameters": 1000, + "method": "ILA", + }, + "gmp": { + **POLYNOMIAL_CONFIGURATIONS["dpd_ila"]["gmp"][ + "basis_configuration" + ], + "real_parameters": 1000, + "method": "ILA", + }, + "gru": {"hidden_size": 16, "real_parameters": 994}, + "tres_gru": {"hidden_size": 15, "real_parameters": 999}, + }, + "dpd_tres_deltagru_pa_surrogate_comparison": { + "model": { + "hidden_size": 15, + "real_parameters": 999, + "delta_thresholds": { + "input": 0.0, + "hidden": 0.0, + }, + }, + "pa_surrogates": { + "tres_gru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": None, + }, + "tres_deltagru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": { + "input": 0.0, + "hidden": 0.0, + }, + }, + }, + }, + "dpd_pa_surrogate": { + "backbone": "tres_gru", + "hidden_size": 27, + "real_parameters": 2751, + }, + "polynomial_solvers": { + "default": { + "implementation": "torch.linalg.lstsq", + "mode": "gels", + "driver": "gels", + "dtype": "torch.complex64", + "column_scaling": "l2", + "regularization": None, + }, + "gmp_pa": { + "implementation": "torch.linalg.svd", + "mode": "truncated_svd", + "driver": "gesvdj", + "dtype": "torch.complex64", + "column_scaling": "l2", + "regularization": { + "type": "truncated_svd", + "relative_cutoff": GMP_PA_SVD_RCOND, + }, + }, + "segment_boundary_policy": ( + "zero delay state at every nperseg boundary" + ), + }, + } + + +def validate_recipe_file(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"Benchmark recipe is missing: {path}") + with path.open() as handle: + recipe = json.load(handle) + require_equal(recipe, expected_runner_recipe(), "recorded benchmark recipe") + return recipe + + +def validate_learning_rate_schedule( + history_rows: list[dict[str, str]], + *, + model_id: str, + initial_learning_rate: float, + minimum_learning_rate: float, + decay_factor: float, +) -> None: + """Validate the observable portion of the configured plateau schedule.""" + learning_rates = [as_float(row, "LR") for row in history_rows] + require_close( + learning_rates[0], + initial_learning_rate, + f"initial learning rate for {model_id}", + tolerance=1e-10, + ) + for epoch, learning_rate in enumerate(learning_rates): + if learning_rate < minimum_learning_rate - 1e-10: + raise ValueError( + f"Learning rate at epoch {epoch} for {model_id} is below " + f"the configured minimum: {learning_rate} < " + f"{minimum_learning_rate}" + ) + if learning_rate > initial_learning_rate + 1e-10: + raise ValueError( + f"Learning rate at epoch {epoch} for {model_id} exceeds " + f"the configured initial value: {learning_rate} > " + f"{initial_learning_rate}" + ) + for epoch, (previous, current) in enumerate( + zip(learning_rates, learning_rates[1:]), start=1 + ): + if current > previous + 1e-10: + raise ValueError( + f"Learning rate increased at epoch {epoch} for {model_id}: " + f"{previous} -> {current}" + ) + if not math.isclose(current, previous, rel_tol=0.0, abs_tol=1e-10): + expected = max(previous * decay_factor, minimum_learning_rate) + require_close( + current, + expected, + f"scheduled learning rate at epoch {epoch} for {model_id}", + tolerance=1e-8, + ) + + +def collect_training_model( + output_dir: Path, + *, + dataset: str, + step: str, + parent: str | None, + model_id: str, + expected_epochs: int, + expected_batch_size: int, + expected_initial_lr: float, + expected_backbone: str, + expected_hidden_size: int, + expected_parameters: int, + selection_metric: str, + expected_delta_thresholds: dict[str, float] | None = None, + expected_pa_parameters: int | None = None, +) -> dict[str, Any]: + base = Path("log") / dataset / step + save_base = Path("save") / dataset / step + if parent is not None: + base /= parent + save_base /= parent + + best_path = base / "best" / f"{model_id}.csv" + history_path = base / "history" / f"{model_id}.csv" + checkpoint_path = save_base / f"{model_id}.pt" + validate_checkpoint(REPO_ROOT / checkpoint_path, expected_parameters) + + best_rows = read_csv(REPO_ROOT / best_path) + history_rows = read_csv(REPO_ROOT / history_path) + require_equal(len(best_rows), 1, f"best-row count for {model_id}") + require_equal( + len(history_rows), expected_epochs, f"history-row count for {model_id}" + ) + require_equal( + [as_int(row, "EPOCH") for row in history_rows], + list(range(expected_epochs)), + f"epoch sequence for {model_id}", + ) + + best = best_rows[0] + require_equal(as_int(best, "N_EPOCH"), expected_epochs, f"epochs for {model_id}") + require_equal( + as_int(best, "BATCH_SIZE"), + expected_batch_size, + f"batch size for {model_id}", + ) + require_equal(as_int(best, "FRAME_LENGTH"), 200, f"frame length for {model_id}") + require_equal(best["BACKBONE"], expected_backbone, f"backbone for {model_id}") + require_equal( + as_int(best, "HIDDEN_SIZE"), + expected_hidden_size, + f"hidden size for {model_id}", + ) + expected_logged_parameters = expected_parameters + if step == "train_dpd": + if expected_pa_parameters is None: + raise ValueError( + f"Expected PA parameter count is required for DPD model {model_id}" + ) + expected_logged_parameters += expected_pa_parameters + require_equal( + as_int(best, "N_PARAM"), + expected_logged_parameters, + f"logged network parameter count for {model_id}", + ) + if expected_delta_thresholds is not None: + for row_index, row in enumerate(history_rows): + require_close( + as_float(row, "THX"), + expected_delta_thresholds["input"], + f"input delta threshold at row {row_index} for {model_id}", + ) + require_close( + as_float(row, "THH"), + expected_delta_thresholds["hidden"], + f"hidden delta threshold at row {row_index} for {model_id}", + ) + require_close( + as_float(best, "THX"), + expected_delta_thresholds["input"], + f"selected input delta threshold for {model_id}", + ) + require_close( + as_float(best, "THH"), + expected_delta_thresholds["hidden"], + f"selected hidden delta threshold for {model_id}", + ) + validate_learning_rate_schedule( + history_rows, + model_id=model_id, + initial_learning_rate=expected_initial_lr, + minimum_learning_rate=NEURAL_RECIPE["minimum_learning_rate"], + decay_factor=NEURAL_RECIPE["decay_factor"], + ) + + best_epoch = as_int(best, "EPOCH") + logged_selection_values = [ + as_float(row, selection_metric) for row in history_rows + ] + minimum_logged_value = min(logged_selection_values) + tied_best_epochs = { + index + for index, value in enumerate(logged_selection_values) + if value == minimum_logged_value + } + if best_epoch not in tied_best_epochs: + raise ValueError( + f"Selected epoch {best_epoch} for {model_id} is not tied for the " + f"minimum logged {selection_metric}; tied epochs are " + f"{sorted(tied_best_epochs)}" + ) + matching_history_rows = [ + row for row in history_rows if as_int(row, "EPOCH") == best_epoch + ] + require_equal( + len(matching_history_rows), 1, f"selected epoch in history for {model_id}" + ) + for key, value in best.items(): + require_equal( + matching_history_rows[0].get(key), + value, + f"best/history value {key} for {model_id}", + ) + + result = { + "selected_epoch_zero_based": best_epoch, + "completed_epochs": len(history_rows), + "selected_learning_rate": as_float(best, "LR"), + "train_loss": as_float(best, "TRAIN_LOSS"), + "metrics": metric_block(best), + "artifacts": { + "checkpoint": archived_artifact(output_dir, checkpoint_path), + "best_csv": archived_artifact(output_dir, best_path), + "history_csv": archived_artifact(output_dir, history_path), + }, + } + if expected_delta_thresholds is not None: + result["delta_thresholds"] = dict(expected_delta_thresholds) + return result + + +def normalize_polynomial_metrics( + values: dict[str, float], +) -> dict[str, float]: + normalized = { + "nmse_db": float(values["NMSE"]), + "evm_db": float(values["EVM"]), + "aclr_left_db": float(values["ACLR_L"]), + "aclr_right_db": float(values["ACLR_R"]), + "aclr_avg_db": float(values["ACLR_AVG"]), + } + for key, value in normalized.items(): + if not math.isfinite(value): + raise ValueError(f"Polynomial metric {key} is not finite: {value}") + require_close( + normalized["aclr_avg_db"], + (normalized["aclr_left_db"] + normalized["aclr_right_db"]) / 2.0, + "polynomial-model ACLR average", + tolerance=1e-10, + ) + return normalized + + +def finite_number( + raw: dict[str, Any], + key: str, + *, + source_path: Path, + positive: bool = False, + nonnegative: bool = False, +) -> float: + try: + value = float(raw[key]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError(f"{key} is missing or non-numeric in {source_path}") from exc + if not math.isfinite(value): + raise ValueError(f"{key} is not finite in {source_path}: {value}") + if positive and value <= 0.0: + raise ValueError(f"{key} must be positive in {source_path}: {value}") + if nonnegative and value < 0.0: + raise ValueError(f"{key} must be nonnegative in {source_path}: {value}") + return value + + +def collect_polynomial_model( + output_dir: Path, + *, + dataset: str, + slug: str, + task: str, + model_name: str, + expected_pa_hash: str | None = None, + expected_device: str | None = None, + expected_target_gain: float | None = None, +) -> dict[str, Any]: + if task not in POLYNOMIAL_CONFIGURATIONS: + raise ValueError(f"Unsupported polynomial benchmark task: {task}") + filename_task = "pa" if task == "pa_modeling" else "dpd" + source_path = ( + output_dir / f"benchmark_report_{slug}_{filename_task}_{model_name}.json" + ) + with source_path.open() as handle: + raw = json.load(handle) + + configuration = POLYNOMIAL_CONFIGURATIONS[task][model_name] + expected_method = ( + "direct_least_squares" + if task == "pa_modeling" + else "indirect_learning_architecture" + ) + uses_truncated_svd = task == "pa_modeling" and model_name == "gmp" + + require_equal( + raw["schema_version"], + POLYNOMIAL_SCHEMA_VERSION, + f"polynomial schema in {source_path}", + ) + require_equal(raw["dataset"], dataset, f"polynomial dataset in {source_path}") + require_equal(raw["task"], task, f"polynomial task in {source_path}") + require_equal(raw["model"], model_name, f"polynomial model in {source_path}") + require_equal(raw["method"], expected_method, f"method in {source_path}") + expected_solver = ( + "torch.linalg.svd" if uses_truncated_svd else "torch.linalg.lstsq" + ) + expected_solver_mode = "truncated_svd" if uses_truncated_svd else "gels" + expected_solver_driver = "gesvdj" if uses_truncated_svd else "gels" + require_equal(raw["solver"], expected_solver, f"polynomial solver in {source_path}") + require_equal( + raw["solver_mode"], + expected_solver_mode, + f"polynomial solver mode in {source_path}", + ) + require_equal( + raw["solver_driver"], + expected_solver_driver, + f"least-squares driver in {source_path}", + ) + require_equal( + raw["dtype"], + "torch.complex64", + f"least-squares dtype in {source_path}", + ) + require_equal( + raw["column_scaling"], + "l2", + f"least-squares column scaling in {source_path}", + ) + if not isinstance(raw["device"], str) or not raw["device"].startswith("cuda:"): + raise ValueError( + f"Polynomial least-squares device must be a CUDA device in " + f"{source_path}: {raw['device']!r}" + ) + require_equal( + raw.get("solver_device"), + raw["device"], + f"least-squares device alias in {source_path}", + ) + if expected_device is not None: + require_equal( + raw["device"], + expected_device, + f"least-squares device in {source_path}", + ) + require_equal( + raw["basis_configuration"], + configuration["basis_configuration"], + f"polynomial basis in {source_path}", + ) + require_equal( + raw["complex_coefficients"], + configuration["complex_coefficients"], + f"least-squares complex coefficient count in {source_path}", + ) + require_equal( + raw["real_parameters"], + configuration["real_parameters"], + f"least-squares parameter count in {source_path}", + ) + require_equal( + raw["parameter_count_convention"], + "two real degrees of freedom per complex coefficient", + f"parameter-count convention in {source_path}", + ) + target_gain = finite_number( + raw, + "target_gain", + source_path=source_path, + positive=True, + ) + if expected_target_gain is not None: + require_close( + target_gain, + expected_target_gain, + f"target gain in {source_path}", + tolerance=1e-6, + ) + rank = raw.get("least_squares_rank") + if uses_truncated_svd: + require_equal( + raw.get("full_rank_assumption"), + False, + f"full-rank assumption for truncated SVD in {source_path}", + ) + require_equal( + raw.get("regularization"), + "truncated_svd", + f"regularization in {source_path}", + ) + svd_rcond = finite_number( + raw, + "svd_rcond", + source_path=source_path, + positive=True, + ) + require_close( + svd_rcond, + GMP_PA_SVD_RCOND, + f"SVD relative cutoff in {source_path}", + tolerance=1e-12, + ) + if not isinstance(rank, int) or isinstance(rank, bool): + raise ValueError( + f"least-squares rank is not an integer in {source_path}: " + f"{rank!r}" + ) + if not 0 < rank < configuration["complex_coefficients"]: + raise ValueError( + f"truncated-SVD rank is outside the expected range in " + f"{source_path}: {rank}" + ) + singular_value_max = finite_number( + raw, + "singular_value_max", + source_path=source_path, + positive=True, + ) + singular_value_min = finite_number( + raw, + "singular_value_min", + source_path=source_path, + positive=True, + ) + singular_value_cutoff = finite_number( + raw, + "singular_value_cutoff", + source_path=source_path, + positive=True, + ) + retained_singular_value_min = finite_number( + raw, + "retained_singular_value_min", + source_path=source_path, + positive=True, + ) + condition_number = finite_number( + raw, + "condition_number", + source_path=source_path, + positive=True, + ) + require_close( + singular_value_cutoff, + singular_value_max * svd_rcond, + f"singular-value cutoff in {source_path}", + tolerance=1e-5, + ) + if not ( + singular_value_min + <= retained_singular_value_min + <= singular_value_max + ): + raise ValueError( + f"retained singular-value bounds are invalid in {source_path}" + ) + if retained_singular_value_min <= singular_value_cutoff: + raise ValueError( + f"retained singular value does not exceed the cutoff in " + f"{source_path}" + ) + if not math.isclose( + condition_number, + singular_value_max / singular_value_min, + rel_tol=1e-6, + ): + raise ValueError( + f"condition number is inconsistent with singular values in " + f"{source_path}" + ) + else: + svd_rcond = None + singular_value_max = None + singular_value_min = None + singular_value_cutoff = None + retained_singular_value_min = None + condition_number = None + require_equal( + raw.get("regularization"), + None, + f"regularization in {source_path}", + ) + require_equal( + raw.get("svd_rcond"), + None, + f"SVD relative cutoff in {source_path}", + ) + for key in ( + "singular_value_max", + "singular_value_min", + "singular_value_cutoff", + "retained_singular_value_min", + "condition_number", + ): + require_equal( + raw.get(key), + None, + f"{key} in {source_path}", + ) + require_equal( + raw.get("full_rank_assumption"), + True, + f"full-rank assumption for CUDA gels in {source_path}", + ) + require_equal(rank, None, f"least-squares rank in {source_path}") + + residual = finite_number( + raw, + "training_relative_residual", + source_path=source_path, + nonnegative=True, + ) + column_scale_min = finite_number( + raw, + "column_scale_min", + source_path=source_path, + positive=True, + ) + column_scale_max = finite_number( + raw, + "column_scale_max", + source_path=source_path, + positive=True, + ) + if column_scale_max < column_scale_min: + raise ValueError( + f"column_scale_max is below column_scale_min in {source_path}: " + f"{column_scale_max} < {column_scale_min}" + ) + coefficient_l2_norm = finite_number( + raw, + "coefficient_l2_norm", + source_path=source_path, + positive=True, + ) + + expected_sample_rate = DATASETS[dataset]["sample_rate_hz"] + expected_nperseg = DATASETS[dataset]["nperseg"] + require_close( + float(raw["sample_rate_hz"]), + expected_sample_rate, + f"sample rate in {source_path}", + ) + require_equal( + raw["nperseg"], + expected_nperseg, + f"segment length in {source_path}", + ) + pa_checkpoint_binding: dict[str, str] | None = None + pa_evaluation_device: str | None = None + if task == "dpd_ila": + if expected_pa_hash is None: + raise ValueError( + f"Expected PA hash is required to collect DPD ILA result: " + f"{source_path}" + ) + expected_checkpoint = ( + Path("save") + / dataset + / "train_pa" + / f"{NEURAL_PA_MODELS[REFERENCE_PA_KEY]['model_id']}.pt" + ).as_posix() + require_equal( + raw["pa_checkpoint"], + expected_checkpoint, + f"PA checkpoint path in {source_path}", + ) + require_equal( + raw["pa_checkpoint_sha256"], + expected_pa_hash, + f"PA checkpoint hash in {source_path}", + ) + pa_evaluation_device = raw.get("pa_evaluation_device") + if expected_device is not None: + require_equal( + pa_evaluation_device, + expected_device, + f"PA evaluation device in {source_path}", + ) + pa_checkpoint_binding = { + "checkpoint": raw["pa_checkpoint"], + "sha256": raw["pa_checkpoint_sha256"], + } + else: + unexpected_binding_fields = { + key + for key in ("pa_checkpoint", "pa_checkpoint_sha256") + if key in raw + } + if unexpected_binding_fields: + raise ValueError( + f"Direct PA polynomial result must not contain a PA checkpoint " + f"binding in {source_path}: {sorted(unexpected_binding_fields)}" + ) + + validation_metrics = normalize_polynomial_metrics(raw["validation"]) + test_metrics = normalize_polynomial_metrics(raw["test"]) + + return { + "task": raw["task"], + "method": raw["method"], + "solver": raw["solver"], + "solver_mode": raw["solver_mode"], + "solver_driver": raw["solver_driver"], + "dtype": raw["dtype"], + "device": raw["device"], + "column_scaling": raw["column_scaling"], + "column_scale_min": column_scale_min, + "column_scale_max": column_scale_max, + "training_relative_residual": residual, + "regularization": raw.get("regularization"), + "svd_rcond": svd_rcond, + "singular_value_max": singular_value_max, + "singular_value_min": singular_value_min, + "singular_value_cutoff": singular_value_cutoff, + "retained_singular_value_min": retained_singular_value_min, + "condition_number": condition_number, + "coefficient_l2_norm": coefficient_l2_norm, + "full_rank_assumption": raw.get("full_rank_assumption"), + "basis_configuration": raw["basis_configuration"], + "complex_coefficients": int(raw["complex_coefficients"]), + "real_parameters": int(raw["real_parameters"]), + "parameter_count_convention": raw["parameter_count_convention"], + "target_gain": target_gain, + "least_squares_rank": int(rank) if rank is not None else None, + "sample_rate_hz": float(raw["sample_rate_hz"]), + "nperseg": int(raw["nperseg"]), + "pa_evaluation_device": pa_evaluation_device, + "pa_checkpoint_binding": pa_checkpoint_binding, + "metrics": { + "validation": validation_metrics, + "test": test_metrics, + }, + "artifact": { + "path": source_path.relative_to(output_dir).as_posix(), + "sha256": sha256_file(source_path), + "size_bytes": source_path.stat().st_size, + }, + } + + +def signal_record(dataset: str) -> dict[str, Any]: + import numpy as np + + arrays = [] + dataset_files: dict[str, str] = {} + train_input: Any = None + train_output: Any = None + for split in ("train", "val", "test"): + input_path = Path("datasets") / dataset / f"{split}_input.csv" + absolute_path = REPO_ROOT / input_path + input_iq = np.loadtxt( + absolute_path, + delimiter=",", + skiprows=1, + dtype=np.float64, + ) + arrays.append(input_iq) + dataset_files[input_path.as_posix()] = sha256_file(absolute_path) + output_path = Path("datasets") / dataset / f"{split}_output.csv" + absolute_output_path = REPO_ROOT / output_path + dataset_files[output_path.as_posix()] = sha256_file(absolute_output_path) + if split == "train": + train_input = input_iq + train_output = np.loadtxt( + absolute_output_path, + delimiter=",", + skiprows=1, + dtype=np.float64, + ) + + iq = np.vstack(arrays) + power = np.square(iq[:, 0]) + np.square(iq[:, 1]) + input_peak = float( + np.max(np.hypot(train_input[:, 0], train_input[:, 1])) + ) + output_peak = float( + np.max(np.hypot(train_output[:, 0], train_output[:, 1])) + ) + mean_power = float(np.mean(power)) + ccdf_probability = 1e-5 + ccdf_power = float( + np.quantile(power, 1.0 - ccdf_probability, method="linear") + ) + return { + "samples": int(power.size), + "target_gain": output_peak / input_peak, + "papr": { + "ccdf_probability": ccdf_probability, + "quantile_method": "numpy linear", + "ccdf_db": float(10.0 * np.log10(ccdf_power / mean_power)), + "absolute_peak_db": float( + 10.0 * np.log10(float(np.max(power)) / mean_power) + ), + }, + "dataset_file_sha256": dataset_files, + } + + +def package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def command_output(command: list[str]) -> str | None: + try: + result = subprocess.run( + command, + cwd=REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() + + +def environment_record(device: int) -> dict[str, Any]: + import torch + + gpu_name = ( + torch.cuda.get_device_name(device) + if torch.cuda.is_available() and device < torch.cuda.device_count() + else None + ) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "numpy": package_version("numpy"), + "scipy": package_version("scipy"), + "pandas": package_version("pandas"), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "gpu": gpu_name, + "gpu_device_index": device, + "training_reproducibility_level": "soft", + "training_deterministic_algorithms": False, + "training_cudnn_benchmark": True, + } + + +def git_record() -> dict[str, Any]: + status = command_output(["git", "status", "--short", "--untracked-files=all"]) + return { + "commit": command_output(["git", "rev-parse", "HEAD"]), + "branch": command_output(["git", "branch", "--show-current"]), + "dirty": bool(status), + "status": status.splitlines() if status else [], + } + + +def runner_start_git_record(output_dir: Path) -> dict[str, Any]: + commit_path = output_dir / "git_commit_before.txt" + branch_path = output_dir / "git_branch_before.txt" + status_path = output_dir / "git_status_before.txt" + if not all(path.is_file() for path in (commit_path, branch_path, status_path)): + raise FileNotFoundError( + "Runner start-state files are missing; use " + "benchmark/reproduce_benchmark_report.sh." + ) + status = status_path.read_text().strip() + return { + "commit": commit_path.read_text().strip(), + "branch": branch_path.read_text().strip(), + "dirty": bool(status), + "status": status.splitlines() if status else [], + } + + +def capture_context( + output_dir: Path, + device: int, + *, + use_runner_start_git: bool, +) -> dict[str, Any]: + return { + "captured_at_utc": datetime.now(timezone.utc).isoformat(), + "git": ( + runner_start_git_record(output_dir) + if use_runner_start_git + else git_record() + ), + "environment": environment_record(device), + "source_file_sha256": { + path: sha256_file(REPO_ROOT / path) for path in source_files() + }, + "datasets": { + dataset: { + "spec_sha256": sha256_file( + REPO_ROOT / "datasets" / dataset / "spec.json" + ), + "signal": signal_record(dataset), + } + for dataset in DATASETS + }, + } + + +def validate_context_stability( + initial: dict[str, Any], + completion: dict[str, Any], +) -> None: + require_equal( + completion["git"]["commit"], + initial["git"]["commit"], + "Git commit across the benchmark run", + ) + require_equal( + completion["git"]["branch"], + initial["git"]["branch"], + "Git branch across the benchmark run", + ) + require_equal( + completion["environment"], + initial["environment"], + "software and hardware environment across the benchmark run", + ) + require_equal( + completion["source_file_sha256"], + initial["source_file_sha256"], + "source files across the benchmark run", + ) + require_equal( + completion["datasets"], + initial["datasets"], + "dataset files and signal evidence across the benchmark run", + ) + + +def load_and_validate_context( + output_dir: Path, + device: int, +) -> tuple[dict[str, Any], dict[str, Any]]: + context_path = output_dir / "context_before.json" + if not context_path.is_file(): + raise FileNotFoundError( + "Pre-run context snapshot is missing; run the collector once with " + "--snapshot-context before training." + ) + with context_path.open() as handle: + initial = json.load(handle) + snapshot_metadata = validate_source_snapshot( + output_dir / "source_snapshot.tar.gz", + initial["source_file_sha256"], + ) + require_equal( + initial.get("source_snapshot"), + snapshot_metadata, + "source snapshot metadata", + ) + completion = capture_context( + output_dir, + device, + use_runner_start_git=False, + ) + validate_context_stability(initial, completion) + return initial, completion + + +def collect_manifest( + output_dir: Path, + device: int, + initial_context: dict[str, Any], + completion_context: dict[str, Any], +) -> dict[str, Any]: + pa_bindings = read_pa_bindings(output_dir) + commands_path = output_dir / "commands.log" + validate_runner_commands( + commands_path, + output_dir=output_dir, + device=device, + ) + recipe_path = output_dir / "recipe.json" + validate_recipe_file(recipe_path) + manifest: dict[str, Any] = { + "schema_version": RUN_SCHEMA_VERSION, + "artifact_kind": "opendpd_benchmark_reproduction_run", + "started_at_utc": initial_context["captured_at_utc"], + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "canonical_report": "benchmark/benchmark_report.md", + "git": initial_context["git"], + "environment": initial_context["environment"], + "completion_context": { + "captured_at_utc": completion_context["captured_at_utc"], + "git_commit": completion_context["git"]["commit"], + "git_branch": completion_context["git"]["branch"], + "validated_unchanged": True, + }, + "run_file_sha256": { + "commands.log": sha256_file(commands_path), + "context_before.json": sha256_file( + output_dir / "context_before.json" + ), + "git_commit_before.txt": sha256_file( + output_dir / "git_commit_before.txt" + ), + "git_branch_before.txt": sha256_file( + output_dir / "git_branch_before.txt" + ), + "git_status_before.txt": sha256_file( + output_dir / "git_status_before.txt" + ), + "pa_checkpoint_bindings.tsv": sha256_file( + output_dir / "pa_checkpoint_bindings.tsv" + ), + "recipe.json": sha256_file(recipe_path), + "source_snapshot.tar.gz": sha256_file( + output_dir / "source_snapshot.tar.gz" + ), + }, + "source_snapshot": initial_context["source_snapshot"], + "source_file_sha256": initial_context["source_file_sha256"], + "recipes": { + "neural_pa": dict(NEURAL_RECIPE), + "neural_dpd": dict(NEURAL_RECIPE), + "neural_dpd_pa_surrogate_sensitivity": { + **NEURAL_RECIPE, + "dpd_model": dict(TRES_DELTAGRU_DPD_MODEL), + "pa_surrogates": [ + REFERENCE_PA_KEY, + ALTERNATE_PA_KEY, + ], + }, + "polynomial_pa": { + "method": "direct_least_squares", + "column_scaling": "l2", + "solvers": { + "mp": { + "implementation": "torch.linalg.lstsq", + "mode": "gels", + "driver": "gels", + "regularization": None, + }, + "gmp": { + "implementation": "torch.linalg.svd", + "mode": "truncated_svd", + "driver": "gesvdj", + "regularization": { + "type": "truncated_svd", + "relative_cutoff": GMP_PA_SVD_RCOND, + }, + }, + }, + "parameter_convention": ( + "two real degrees of freedom per complex coefficient" + ), + "models": POLYNOMIAL_CONFIGURATIONS["pa_modeling"], + }, + "polynomial_dpd": { + "method": "indirect_learning_architecture", + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "column_scaling": "l2", + "parameter_convention": ( + "two real degrees of freedom per complex coefficient" + ), + "models": POLYNOMIAL_CONFIGURATIONS["dpd_ila"], + }, + }, + "metric_note": ( + "EVM is the repository-specific normalized mean absolute complex-" + "spectrum error across configured main-channel subchannels, not " + "demodulated constellation EVM. DPD metrics are simulated through " + "a frozen learned PA surrogate." + ), + "datasets": {}, + } + + for dataset, configuration in DATASETS.items(): + pa_models: dict[str, Any] = {} + for model_name in ("mp", "gmp"): + result = collect_polynomial_model( + output_dir, + dataset=dataset, + slug=configuration["slug"], + task="pa_modeling", + model_name=model_name, + expected_device=f"cuda:{device}", + expected_target_gain=initial_context["datasets"][dataset][ + "signal" + ]["target_gain"], + ) + result["model"] = { + "display_name": model_name.upper(), + "backbone": model_name, + "parameters": result["real_parameters"], + "parameter_convention": ( + "two real degrees of freedom per complex coefficient" + ), + } + pa_models[model_name] = result + + for key, model_configuration in NEURAL_PA_MODELS.items(): + result = collect_training_model( + output_dir, + dataset=dataset, + step="train_pa", + parent=None, + model_id=model_configuration["model_id"], + expected_epochs=NEURAL_RECIPE["epochs"], + expected_batch_size=NEURAL_RECIPE["batch_size"], + expected_initial_lr=NEURAL_RECIPE["initial_learning_rate"], + expected_backbone=model_configuration["backbone"], + expected_hidden_size=model_configuration["hidden_size"], + expected_parameters=model_configuration["parameters"], + selection_metric="VAL_NMSE", + expected_delta_thresholds=model_configuration.get( + "delta_thresholds" + ), + ) + result["task"] = "pa_modeling" + result["method"] = "supervised_training" + result["model"] = dict(model_configuration) + pa_models[key] = result + + reference_pa = pa_models[REFERENCE_PA_KEY] + pa_hash = reference_pa["artifacts"]["checkpoint"]["sha256"] + expected_pa_checkpoint = ( + Path("save") + / dataset + / "train_pa" + / f"{NEURAL_PA_MODELS[REFERENCE_PA_KEY]['model_id']}.pt" + ).as_posix() + + dpd_models: dict[str, Any] = {} + for model_name in ("mp", "gmp"): + result = collect_polynomial_model( + output_dir, + dataset=dataset, + slug=configuration["slug"], + task="dpd_ila", + model_name=model_name, + expected_pa_hash=pa_hash, + expected_device=f"cuda:{device}", + expected_target_gain=initial_context["datasets"][dataset][ + "signal" + ]["target_gain"], + ) + result["model"] = { + "display_name": model_name.upper(), + "backbone": model_name, + "parameters": result["real_parameters"], + "parameter_convention": ( + "two real degrees of freedom per complex coefficient" + ), + } + dpd_models[model_name] = result + + for key, model_configuration in NEURAL_DPD_MODELS.items(): + result = collect_training_model( + output_dir, + dataset=dataset, + step="train_dpd", + parent=REFERENCE_PA_PARENT, + model_id=model_configuration["model_id"], + expected_epochs=NEURAL_RECIPE["epochs"], + expected_batch_size=NEURAL_RECIPE["batch_size"], + expected_initial_lr=NEURAL_RECIPE["initial_learning_rate"], + expected_backbone=model_configuration["backbone"], + expected_hidden_size=model_configuration["hidden_size"], + expected_parameters=model_configuration["parameters"], + selection_metric="VAL_ACLR_AVG", + expected_delta_thresholds=model_configuration.get( + "delta_thresholds" + ), + expected_pa_parameters=reference_pa["model"]["parameters"], + ) + result["task"] = "dpd_training" + result["method"] = "direct_learning_architecture" + result["model"] = dict(model_configuration) + binding = pa_bindings.get((dataset, key, REFERENCE_PA_KEY)) + if binding is None: + raise ValueError( + f"Missing PA checkpoint binding for " + f"{dataset}/{key}/{REFERENCE_PA_KEY}" + ) + require_equal( + binding["checkpoint"], + expected_pa_checkpoint, + f"PA checkpoint path binding for {dataset}/{key}", + ) + require_equal( + binding["sha256"], + pa_hash, + f"PA checkpoint hash binding for {dataset}/{key}", + ) + result["pa_checkpoint_binding"] = binding + dpd_models[key] = result + + delta_dpd_by_pa_surrogate: dict[str, Any] = {} + for pa_model_key, pa_parent in ( + (REFERENCE_PA_KEY, REFERENCE_PA_PARENT), + (ALTERNATE_PA_KEY, ALTERNATE_PA_PARENT), + ): + pa_result = pa_models[pa_model_key] + pa_checkpoint = ( + Path("save") + / dataset + / "train_pa" + / f"{NEURAL_PA_MODELS[pa_model_key]['model_id']}.pt" + ).as_posix() + result = collect_training_model( + output_dir, + dataset=dataset, + step="train_dpd", + parent=pa_parent, + model_id=TRES_DELTAGRU_DPD_MODEL["model_id"], + expected_epochs=NEURAL_RECIPE["epochs"], + expected_batch_size=NEURAL_RECIPE["batch_size"], + expected_initial_lr=NEURAL_RECIPE["initial_learning_rate"], + expected_backbone=TRES_DELTAGRU_DPD_MODEL["backbone"], + expected_hidden_size=TRES_DELTAGRU_DPD_MODEL["hidden_size"], + expected_parameters=TRES_DELTAGRU_DPD_MODEL["parameters"], + selection_metric="VAL_ACLR_AVG", + expected_delta_thresholds=TRES_DELTAGRU_DPD_MODEL[ + "delta_thresholds" + ], + expected_pa_parameters=pa_result["model"]["parameters"], + ) + result["task"] = "dpd_training" + result["method"] = "direct_learning_architecture" + result["model"] = dict(TRES_DELTAGRU_DPD_MODEL) + binding = pa_bindings.get( + (dataset, "tres_deltagru", pa_model_key) + ) + if binding is None: + raise ValueError( + f"Missing PA checkpoint binding for " + f"{dataset}/tres_deltagru/{pa_model_key}" + ) + require_equal( + binding["checkpoint"], + pa_checkpoint, + ( + "PA checkpoint path binding for " + f"{dataset}/tres_deltagru/{pa_model_key}" + ), + ) + require_equal( + binding["sha256"], + pa_result["artifacts"]["checkpoint"]["sha256"], + ( + "PA checkpoint hash binding for " + f"{dataset}/tres_deltagru/{pa_model_key}" + ), + ) + result["pa_checkpoint_binding"] = binding + result["pa_surrogate"] = { + "model_key": pa_model_key, + "model": dict(NEURAL_PA_MODELS[pa_model_key]), + "checkpoint": pa_checkpoint, + "checkpoint_sha256": binding["sha256"], + "selection_metric": "validation NMSE", + } + delta_dpd_by_pa_surrogate[pa_model_key] = result + + manifest["datasets"][dataset] = { + "spec_sha256": initial_context["datasets"][dataset]["spec_sha256"], + "signal": initial_context["datasets"][dataset]["signal"], + "reference_pa": { + "model_key": REFERENCE_PA_KEY, + "checkpoint": expected_pa_checkpoint, + "checkpoint_sha256": pa_hash, + "selection_metric": "validation NMSE", + }, + "pa_models": pa_models, + "dpd_models": dpd_models, + "dpd_tres_deltagru_by_pa_surrogate": ( + delta_dpd_by_pa_surrogate + ), + } + + return manifest + + +def publish_manifest_from_run( + manifest: dict[str, Any], + run_directory: Path, +) -> dict[str, Any]: + """Bind post-collection source/job evidence into a canonical manifest.""" + source_schema_version = manifest["schema_version"] + source_result_count = sum( + len(dataset["pa_models"]) + + len(dataset["dpd_models"]) + + len(dataset["dpd_tres_deltagru_by_pa_surrogate"]) + for dataset in manifest["datasets"].values() + ) + if source_schema_version != RUN_SCHEMA_VERSION: + raise ValueError( + "Unsupported source manifest schema for publication: " + f"{source_schema_version}" + ) + context_path = run_directory / "context_before.json" + original_manifest_path = run_directory / "benchmark_report_results.json" + jobs_path = run_directory / "jobs.tsv" + recipe_path = run_directory / "recipe.json" + commands_path = run_directory / "commands.log" + with context_path.open() as handle: + context = json.load(handle) + + snapshot_path = run_directory / "source_snapshot.tar.gz" + snapshot_metadata = validate_source_snapshot( + snapshot_path, + context["source_file_sha256"], + ) + snapshot_time = datetime.fromtimestamp( + snapshot_path.stat().st_mtime, + timezone.utc, + ) + started = datetime.fromisoformat(manifest["started_at_utc"]) + collected = datetime.fromisoformat(manifest["generated_at_utc"]) + if not started <= snapshot_time <= collected: + raise ValueError( + "Source snapshot was not created during the recorded benchmark run: " + f"{snapshot_time.isoformat()} is outside " + f"{started.isoformat()} .. {collected.isoformat()}" + ) + + with jobs_path.open(newline="") as handle: + jobs = list(csv.DictReader(handle, delimiter="\t")) + command_labels = set(read_recorded_commands(commands_path)) + job_labels = [row["label"] for row in jobs] + require_equal(len(job_labels), len(set(job_labels)), "unique job labels") + require_equal(set(job_labels), command_labels, "completed benchmark jobs") + for row in jobs: + require_equal( + int(row["exit_status"]), + 0, + f"exit status for job {row['label']}", + ) + + published = copy.deepcopy(manifest) + published["schema_version"] = RUN_SCHEMA_VERSION + for dataset in published["datasets"].values(): + dataset["pa_models"] = { + key: dataset["pa_models"][key] for key in PA_MODEL_ORDER + } + dataset["dpd_models"] = { + key: dataset["dpd_models"][key] for key in DPD_MODEL_ORDER + } + dataset["dpd_tres_deltagru_by_pa_surrogate"] = { + key: dataset["dpd_tres_deltagru_by_pa_surrogate"][key] + for key in (REFERENCE_PA_KEY, ALTERNATE_PA_KEY) + } + published["source_snapshot"] = snapshot_metadata + published["run_file_sha256"].update( + { + "jobs.tsv": sha256_file(jobs_path), + "recipe.json": sha256_file(recipe_path), + "source_snapshot.tar.gz": sha256_file(snapshot_path), + } + ) + for stage in ("neural_pa", "neural_dpd"): + published["recipes"][stage].update(copy.deepcopy(NEURAL_RECIPE)) + published["metric_note"] = ( + "EVM is the repository-specific normalized mean absolute complex-" + "spectrum error across configured main-channel subchannels, not " + "demodulated constellation EVM. DPD metrics are simulated through a " + "frozen learned PA surrogate." + ) + published["publication"] = { + "source_run_id": run_directory.name, + "original_schema_version": source_schema_version, + "original_collector_manifest_sha256": sha256_file( + original_manifest_path + ), + "source_snapshot_created_during_run_utc": snapshot_time.isoformat(), + "source_snapshot_member_hashes_match_pre_run_context": True, + "source_job_count": len(jobs), + "source_result_count": source_result_count, + "published_result_count": ( + len(published["datasets"]) + * (len(PA_MODEL_ORDER) + len(DPD_MODEL_ORDER) + 2) + ), + "published_models": list(PA_MODEL_ORDER), + "published_models_by_stage": { + "pa_modeling": list(PA_MODEL_ORDER), + "dpd": list(DPD_MODEL_ORDER), + "dpd_pa_surrogate_sensitivity": ["tres_deltagru"], + }, + "published_dpd_pa_surrogates": [ + REFERENCE_PA_KEY, + ALTERNATE_PA_KEY, + ], + "all_jobs_exit_zero": True, + "total_job_duration_seconds": sum( + int(row["duration_seconds"]) for row in jobs + ), + "note": ( + "The canonical manifest contains the published model subset. " + "Source-run hashes, job totals, and archives cover the complete " + "source run and can include experiments outside that subset; " + "source-archive members were verified against the pre-run source " + "hash map." + ), + } + return published + + +def pair(metrics: dict[str, dict[str, float]], key: str) -> str: + return ( + f"{metrics['validation'][key]:.4f} / " + f"{metrics['test'][key]:.4f}" + ) + + +def aclr_triplet(metrics: dict[str, float]) -> str: + return ( + f"{metrics['aclr_left_db']:.4f} / " + f"{metrics['aclr_right_db']:.4f} / " + f"{metrics['aclr_avg_db']:.4f}" + ) + + +def result_method(result: dict[str, Any]) -> str: + if result["method"] == "supervised_training": + return "Supervised, AdamW" + if result["method"] == "direct_learning_architecture": + return "DLA, AdamW" + if result["method"] == "direct_least_squares": + if result.get("regularization") == "truncated_svd": + return ( + "Truncated SVD " + f"(rank {result['least_squares_rank']:,}/" + f"{result['complex_coefficients']:,})" + ) + return "Direct least squares" + if result["method"] == "indirect_learning_architecture": + return "ILA, least squares" + raise ValueError(f"Unknown benchmark method: {result['method']}") + + +def selected_epoch(result: dict[str, Any]) -> str: + epoch = result.get("selected_epoch_zero_based") + return "N/A" if epoch is None else str(epoch) + + +def result_table( + results: dict[str, dict[str, Any]], + model_order: tuple[str, ...], +) -> list[str]: + lines = [ + "| Model | Parameters | Method | Selected epoch | " + "NMSE, val / test (dB) | EVM, val / test (dB) | " + "Validation ACLR L / R / avg (dB) | " + "Test ACLR L / R / avg (dB) |", + "|---|---:|---|---:|---:|---:|---:|---:|", + ] + for key in model_order: + result = results[key] + lines.append( + f"| {result['model']['display_name']} | " + f"{result['model']['parameters']:,} | " + f"{result_method(result)} | {selected_epoch(result)} | " + f"{pair(result['metrics'], 'nmse_db')} | " + f"{pair(result['metrics'], 'evm_db')} | " + f"{aclr_triplet(result['metrics']['validation'])} | " + f"{aclr_triplet(result['metrics']['test'])} |" + ) + return lines + + +def delta_dpd_surrogate_table(values: dict[str, Any]) -> list[str]: + lines = [ + "| PA surrogate | PA NMSE, val / test (dB) | DPD parameters | " + "Selected epoch | DPD NMSE, val / test (dB) | " + "DPD EVM, val / test (dB) | DPD ACLR avg, val / test (dB) |", + "|---|---:|---:|---:|---:|---:|---:|", + ] + results = values["dpd_tres_deltagru_by_pa_surrogate"] + for pa_model_key in (REFERENCE_PA_KEY, ALTERNATE_PA_KEY): + result = results[pa_model_key] + pa_result = values["pa_models"][pa_model_key] + lines.append( + f"| {pa_result['model']['display_name']} | " + f"{pair(pa_result['metrics'], 'nmse_db')} | " + f"{result['model']['parameters']:,} | " + f"{selected_epoch(result)} | " + f"{pair(result['metrics'], 'nmse_db')} | " + f"{pair(result['metrics'], 'evm_db')} | " + f"{pair(result['metrics'], 'aclr_avg_db')} |" + ) + return lines + + +def write_results_figure(manifest: dict[str, Any], path: Path) -> None: + """Render a compact test-set comparison directly from the manifest.""" + import matplotlib + + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + colors = { + "mp": "#4477AA", + "gmp": "#66CCEE", + "gru": "#228833", + "tres_gru": "#EE6677", + "tres_deltagru": "#AA3377", + } + panels = ( + ( + "APA_200MHz", + "pa_models", + "nmse_db", + "PA test NMSE (dB)", + PA_MODEL_ORDER, + ), + ( + "DPA_160MHz", + "pa_models", + "nmse_db", + "PA test NMSE (dB)", + PA_MODEL_ORDER, + ), + ( + "APA_200MHz", + "dpd_models", + "aclr_avg_db", + "DPD test ACLR average (dB)", + DPD_MODEL_ORDER, + ), + ( + "DPA_160MHz", + "dpd_models", + "aclr_avg_db", + "DPD test ACLR average (dB)", + DPD_MODEL_ORDER, + ), + ) + + figure, axes = plt.subplots( + 2, + 2, + figsize=(13.5, 9.0), + constrained_layout=True, + ) + figure.suptitle("OpenDPD benchmark — test split", fontsize=17, weight="bold") + for axis, (dataset, result_key, metric, title, model_order) in zip( + axes.flat, + panels, + ): + results = manifest["datasets"][dataset][result_key] + values = [ + results[key]["metrics"]["test"][metric] for key in model_order + ] + validation_winner = min( + model_order, + key=lambda key: results[key]["metrics"]["validation"][metric], + ) + span = max(values) - min(values) + label_offset = max(0.22, span * 0.025) + left_padding = max(0.4, span * 0.035) + right_padding = max(1.0, span * 0.16) + for y_position, (key, value) in enumerate(zip(model_order, values)): + axis.scatter( + value, + y_position, + s=115, + color=colors[key], + edgecolor="black" if key == validation_winner else "white", + linewidth=1.8 if key == validation_winner else 0.8, + zorder=3, + ) + axis.text( + value + label_offset, + y_position, + f"{value:.2f}", + va="center", + ha="left", + fontsize=9, + ) + axis.set_xlim( + min(values) - left_padding, + max(values) + right_padding, + ) + if min(values) < 0.0 < max(values): + axis.axvline(0.0, color="#777777", linestyle="--", linewidth=1) + axis.set_yticks( + range(len(model_order)), + [results[key]["model"]["display_name"] for key in model_order], + ) + axis.invert_yaxis() + axis.grid(axis="x", color="#D8D8D8", linewidth=0.8) + axis.set_axisbelow(True) + axis.set_title(f"{dataset}\n{title}", fontsize=12) + axis.set_xlabel("More negative is better ←") + for spine in ("top", "right"): + axis.spines[spine].set_visible(False) + + figure.text( + 0.5, + -0.01, + "Outlined markers are validation leaders. DPD is simulated through " + "the dataset-specific validation-selected TRes-GRU-H27 PA surrogate.", + ha="center", + fontsize=10, + ) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".png", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + figure.savefig( + temporary_path, + dpi=160, + facecolor="white", + bbox_inches="tight", + ) + os.replace(temporary_path, path) + finally: + plt.close(figure) + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() + + +def write_delta_dpd_figure(manifest: dict[str, Any], path: Path) -> None: + """Render the TRes-DeltaGRU DPD PA-surrogate sensitivity comparison.""" + import matplotlib + + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + pa_order = (REFERENCE_PA_KEY, ALTERNATE_PA_KEY) + colors = { + REFERENCE_PA_KEY: "#EE6677", + ALTERNATE_PA_KEY: "#AA3377", + } + figure, axes = plt.subplots( + 1, + 2, + figsize=(13.5, 4.8), + constrained_layout=True, + ) + figure.suptitle( + "TRes-DeltaGRU-H15 DPD — PA surrogate sensitivity (test split)", + fontsize=16, + weight="bold", + ) + for axis, dataset in zip(axes, DATASETS): + dataset_values = manifest["datasets"][dataset] + results = dataset_values["dpd_tres_deltagru_by_pa_surrogate"] + values = [ + results[key]["metrics"]["test"]["aclr_avg_db"] for key in pa_order + ] + validation_winner = min( + pa_order, + key=lambda key: results[key]["metrics"]["validation"]["aclr_avg_db"], + ) + span = max(values) - min(values) + label_offset = max(0.12, span * 0.08) + for y_position, (key, value) in enumerate(zip(pa_order, values)): + axis.scatter( + value, + y_position, + s=125, + color=colors[key], + edgecolor="black" if key == validation_winner else "white", + linewidth=1.8 if key == validation_winner else 0.8, + zorder=3, + ) + axis.text( + value + label_offset, + y_position, + f"{value:.2f}", + va="center", + ha="left", + fontsize=10, + ) + left_padding = max(0.35, span * 0.25) + right_padding = max(0.8, span * 0.55) + axis.set_xlim(min(values) - left_padding, max(values) + right_padding) + axis.set_yticks( + range(len(pa_order)), + [ + f"{dataset_values['pa_models'][key]['model']['display_name']} PA" + for key in pa_order + ], + ) + axis.invert_yaxis() + axis.grid(axis="x", color="#D8D8D8", linewidth=0.8) + axis.set_axisbelow(True) + axis.set_title( + f"{dataset}\nDPD test ACLR average (dB)", + fontsize=12, + ) + axis.set_xlabel("More negative is better ←") + for spine in ("top", "right"): + axis.spines[spine].set_visible(False) + + figure.text( + 0.5, + -0.02, + "Each point is an independently trained TRes-DeltaGRU-H15 DPD. " + "Outlined markers are selected by validation ACLR average.", + ha="center", + fontsize=10, + ) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".png", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + figure.savefig( + temporary_path, + dpi=160, + facecolor="white", + bbox_inches="tight", + ) + os.replace(temporary_path, path) + finally: + plt.close(figure) + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() + + +def bind_derived_artifacts( + manifest: dict[str, Any], + *, + report_path: Path, + figure_path: Path, + delta_dpd_figure_path: Path, +) -> None: + manifest["derived_artifacts"] = { + report_path.name: { + "sha256": sha256_file(report_path), + "size_bytes": report_path.stat().st_size, + }, + figure_path.name: { + "sha256": sha256_file(figure_path), + "size_bytes": figure_path.stat().st_size, + }, + delta_dpd_figure_path.name: { + "sha256": sha256_file(delta_dpd_figure_path), + "size_bytes": delta_dpd_figure_path.stat().st_size, + }, + } + + +def render_markdown( + manifest: dict[str, Any], + *, + canonical: bool = False, +) -> str: + lines = [ + "# OpenDPD PA Modeling and DPD Benchmark", + "", + "## Technical summary", + "", + "This benchmark evaluates MP, GMP, GRU, TRes-GRU, and " + "TRes-DeltaGRU (THX=THH=0) for PA modeling on APA_200MHz and " + "DPA_160MHz. The DPD comparison evaluates MP, GMP, GRU, and TRes-GRU. " + "PA models use approximately 2,700 real parameters; DPD models use " + "approximately 1,000. Every neural run uses the same 300-epoch " + "optimization recipe. MP PA uses direct least squares and GMP PA uses " + "rank-controlled truncated SVD; their predistorters use the indirect " + "learning architecture (ILA). The four-model DPD comparison uses the " + "dataset-specific " + "TRes-GRU-H27 PA checkpoint selected by validation NMSE. A separate " + "sensitivity experiment trains TRes-DeltaGRU-H15 DPD independently " + "through both the TRes-GRU-H27 and zero-threshold " + "TRes-DeltaGRU-H27 PA surrogates.", + "", + "## Key findings", + "", + ] + + for dataset, values in manifest["datasets"].items(): + best_pa = min( + values["pa_models"].values(), + key=lambda result: result["metrics"]["validation"]["nmse_db"], + ) + best_dpd = min( + values["dpd_models"].values(), + key=lambda result: result["metrics"]["validation"]["aclr_avg_db"], + ) + delta_dpd_results = values["dpd_tres_deltagru_by_pa_surrogate"] + best_delta_pa_key = min( + delta_dpd_results, + key=lambda key: delta_dpd_results[key]["metrics"]["validation"][ + "aclr_avg_db" + ], + ) + best_delta_dpd = delta_dpd_results[best_delta_pa_key] + best_delta_pa = values["pa_models"][best_delta_pa_key] + lines.append( + f"- **{dataset}:** the validation PA leader is " + f"{best_pa['model']['display_name']} at " + f"{best_pa['metrics']['validation']['nmse_db']:.4f} dB NMSE " + f"({best_pa['metrics']['test']['nmse_db']:.4f} dB test). The " + f"validation DPD leader is {best_dpd['model']['display_name']} at " + f"{best_dpd['metrics']['validation']['aclr_avg_db']:.4f} dB ACLR " + f"average ({best_dpd['metrics']['test']['aclr_avg_db']:.4f} dB test). " + f"For TRes-DeltaGRU-H15 DPD, the better validation result uses the " + f"{best_delta_pa['model']['display_name']} PA surrogate at " + f"{best_delta_dpd['metrics']['validation']['aclr_avg_db']:.4f} dB " + f"ACLR average " + f"({best_delta_dpd['metrics']['test']['aclr_avg_db']:.4f} dB test)." + ) + + apa_gmp = manifest["datasets"]["APA_200MHz"]["pa_models"]["gmp"] + lines.extend( + [ + ( + "- **APA GMP stability:** column-normalized truncated SVD at " + f"`rcond={apa_gmp['svd_rcond']:.0e}` retains " + f"{apa_gmp['least_squares_rank']:,}/" + f"{apa_gmp['complex_coefficients']:,} singular directions. " + "Validation/test NMSE is " + f"{apa_gmp['metrics']['validation']['nmse_db']:.2f}/" + f"{apa_gmp['metrics']['test']['nmse_db']:.2f} dB. The fixed " + "cutoff suppresses ill-conditioned delayed-envelope directions " + "instead of applying CUDA `gels`'s invalid full-rank " + "assumption." + ), + "", + "![Test-set PA modeling and DPD results](benchmark_results.png)", + "", + "*Test split; more negative is better. Outlined points are the " + "models selected by validation. DPD is simulated through the " + "dataset-specific TRes-GRU-H27 PA surrogate.*", + "", + "![TRes-DeltaGRU DPD PA-surrogate sensitivity]" + "(benchmark_delta_dpd_results.png)", + "", + "*Test split; each point is an independently trained " + "TRes-DeltaGRU-H15 DPD using the named frozen PA surrogate. " + "These results compare surrogate sensitivity, not measurements " + "from one shared physical PA.*", + ] + ) + + for dataset, values in manifest["datasets"].items(): + lines.extend( + [ + "", + f"## {dataset}", + "", + "### PA modeling", + "", + ] + ) + lines.extend(result_table(values["pa_models"], PA_MODEL_ORDER)) + lines.extend( + [ + "", + "### DPD", + "", + ] + ) + lines.extend(result_table(values["dpd_models"], DPD_MODEL_ORDER)) + lines.extend( + [ + "", + "### TRes-DeltaGRU DPD by PA surrogate", + "", + "Both rows use TRes-DeltaGRU-H15 with 999 parameters and " + "THX=THH=0. They are separately trained through the indicated " + "frozen PA surrogate.", + "", + ] + ) + lines.extend(delta_dpd_surrogate_table(values)) + + lines.extend( + [ + "", + "## Definitions", + "", + "- **NMSE:** normalized mean-square error in dB; more negative is " + "better. The implementation averages per-segment dB ratios rather " + "than pooling all samples into one ratio.", + "- **EVM:** the repository-specific mean absolute complex-spectrum " + "error within the configured main channel, normalized within each " + "subchannel by reference-spectrum magnitude and converted with " + "`20 log10`. It is not demodulated constellation EVM; more negative " + "is better.", + "- **ACLR L / R / avg:** adjacent-subchannel power normalized by the " + "strongest configured main-channel subchannel, plus the arithmetic " + "mean of left and right, in dB. More negative is better.", + "- **Parameters:** neural checkpoint tensor elements. MP/GMP count " + "two real degrees of freedom for every complex coefficient.", + "- **Selected epoch:** zero-based neural checkpoint epoch. It is N/A " + "for closed-form least-squares fits.", + "", + "## Model configurations", + "", + "| Stage | MP | GMP | GRU | TRes-GRU | TRes-DeltaGRU |", + "|---|---|---|---|---|---|", + "| PA modeling | K=9, Q=150 | Ka/La=5/30; " + "Kb/Lb/Mb=4/30/5; Kc/Lc/Mc=4/30/5 | H28 | H27 | " + "H27, THX=THH=0 |", + "| DPD | K=5, Q=100 | Ka/La=5/20; " + "Kb/Lb/Mb=4/20/3; Kc/Lc/Mc=4/20/2 | H16 | H15 | " + "H15, THX=THH=0 (PA-surrogate sensitivity) |", + "", + "## Temporal context and sequence boundaries", + "", + "In sequence interiors, PA GMP uses up to five future samples and " + "DPD GMP uses up to two through their leading-envelope terms. " + "TRes-GRU and TRes-DeltaGRU use one-sample right context in their " + "recurrent features and 16-sample right context in their dilated " + "residual convolution. MP and GRU use no explicit future samples.", + "", + "These are offline segmented evaluations. GMP delay accesses are " + "zero-filled and reset at each `nperseg` boundary. In both TRes " + "models, `torch.roll(..., shifts=-1)` wraps the final position to " + "the first sample of the same supplied sequence; the convolution " + "zero-pads both boundaries, and recurrent state resets for each " + "sequence. Neural optimization uses overlapping 200-sample frames " + "with stride 1, while validation and test use independent " + "`nperseg` segments.", + "", + "## Methodology", + "", + "Neural PA and DPD models use batch size 64, 300 epochs, AdamW with " + "MSE loss, initial learning rate 5e-3, and ReduceLROnPlateau with " + "factor 0.5, patience 5, and minimum learning rate 5e-5. Frame " + "length is 200, frame stride is 1, and seed is 0. PA checkpoints " + "are selected by minimum validation NMSE; neural DPD checkpoints " + "are selected by minimum validation ACLR average.", + "", + "The TRes-DeltaGRU PA and DPD runs use input-delta threshold THX=0 " + "and hidden-state-delta threshold THH=0. This disables " + "threshold-induced temporal pruning; exact arithmetic deltas may " + "still naturally be zero. The configuration therefore evaluates " + "the dense zero-threshold recurrence, not a sparsity or efficiency " + "claim.", + "", + "AdamW uses weight decay 0.01, betas (0.9, 0.999), and epsilon " + "1e-8. The scheduler uses relative threshold 1e-4, cooldown 0, and " + "epsilon 1e-8.", + "", + "MP and GMP are complex polynomial models fit after L2 column " + "scaling. MP and both ILA-DPD fits use `torch.linalg.lstsq` " + "(`gels`). GMP PA modeling uses `torch.linalg.svd` (`gesvdj`) " + f"with a fixed relative cutoff of {GMP_PA_SVD_RCOND:.0e}; the " + "effective retained rank is reported with each GMP PA result. " + "These closed-form fits do not use the neural batch, epoch, " + "optimizer, or learning-rate settings. PA polynomial fits map " + "measured PA input to output directly. DPD polynomial fits use " + "ILA, fitting a postdistorter and copying its coefficients to the " + "predistorter. No ridge penalty or validation-tuned regularization " + "is applied.", + "", + "Each dataset uses its independently trained TRes-GRU-H27 PA " + "surrogate for the four-model DPD comparison. The two " + "TRes-DeltaGRU-H15 sensitivity rows are separate DPD training runs, " + "one through that TRes-GRU PA and one through the independently " + "trained TRes-DeltaGRU-H27 PA. Training data supply all gradient " + "and least-squares fits. Neural validation metrics drive the " + "learning-rate schedule and checkpoint selection; test data are " + "not used for fitting, scheduling, or selection.", + "", + "## Limitations and robustness", + "", + "- Neural execution uses soft reproducibility with cuDNN benchmark " + "enabled, so repeated runs can differ slightly.", + "- DPD scores measure simulated performance through a learned PA " + "surrogate, not a fresh over-the-air or bench measurement.", + "- One seed is evaluated. The reported table is not a distribution " + "over training runs.", + "- The five PA candidates and four primary DPD candidates are " + "matched approximately by real parameter count, not by FLOPs, " + "latency, memory traffic, or fit time. The PA-surrogate sensitivity " + "experiment adds two independently trained TRes-DeltaGRU DPD runs " + "per dataset.", + "- Results obtained through different learned PA surrogates are " + "simulator-sensitivity evidence; they are not a controlled ranking " + "against one shared physical PA response.", + "- CUDA `gels` assumes a full-rank design matrix and does not " + "return numerical rank. It remains in use for MP and ILA-DPD; the " + "GMP PA path records its SVD spectrum, cutoff, and retained rank.", + "- GMP PA has 1,350 stored complex coefficients (2,700 nominal real " + "parameters), but truncated SVD reduces its effective rank. The " + "comparison is matched by stored coefficient count, not effective " + "degrees of freedom.", + "- Look-ahead is an input dependency, not measured inference " + "latency. Streaming reformulations and continuous boundary/state " + "handling are not evaluated.", + "", + "## Provenance", + "", + f"- Generated: `{manifest['generated_at_utc']}`", + f"- Git commit: `{manifest['git']['commit']}`", + f"- Git branch: `{manifest['git']['branch']}`", + f"- Git working tree at launch: " + f"`{'dirty' if manifest['git']['dirty'] else 'clean'}`; exact " + "source hashes and start status are retained in the machine " + "evidence.", + f"- Python: `{manifest['environment']['python']}`", + f"- PyTorch: `{manifest['environment']['torch']}`", + f"- CUDA: `{manifest['environment']['cuda']}`", + f"- GPU: `{manifest['environment']['gpu']}`", + ] + ) + if canonical: + publication = manifest.get("publication") + if publication is not None: + lines.append( + "- Canonical evidence schema: " + f"`{manifest['schema_version']}` (published from collector " + f"schema `{publication['original_schema_version']}`); the " + "source archive and completed job ledger are cryptographically " + "bound in the machine-readable evidence." + ) + lines.extend( + [ + "- Reproduce the full matrix: " + "[`reproduce_benchmark_report.sh`](reproduce_benchmark_report.sh)", + "- Machine-readable evidence: " + "[`results/benchmark_report_results.json`]" + "(results/benchmark_report_results.json)", + "- Each reproduction run writes exact commands, a verified " + "source snapshot, raw polynomial results, checkpoints, and CSV " + "logs to its timestamped evidence directory.", + ] + ) + else: + lines.extend( + [ + "- Exact launch commands: [`commands.log`](commands.log)", + "- Machine-readable evidence: " + "[`benchmark_report_results.json`](benchmark_report_results.json)", + "- Source snapshot: " + "[`source_snapshot.tar.gz`](source_snapshot.tar.gz)", + "- Archived checkpoints and CSV logs: [`artifacts/`](artifacts/)", + ] + ) + lines.append("") + return "\n".join(lines) + + +def write_atomic(path: Path, content: str) -> None: + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + handle.write(content) + temporary_path = Path(handle.name) + os.replace(temporary_path, path) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--output-dir", + type=Path, + required=True, + help="Run directory containing raw least-squares JSON and archived artifacts.", + ) + parser.add_argument( + "--device", + type=int, + default=0, + help="CUDA device index used by the benchmark run.", + ) + parser.add_argument( + "--snapshot-context", + action="store_true", + help="Capture pre-run source, dataset, Git, and environment evidence.", + ) + parser.add_argument( + "--overwrite-output", + action="store_true", + help="Replace an existing collector snapshot or final manifest.", + ) + args = parser.parse_args() + output_dir = args.output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + if args.snapshot_context: + context_path = output_dir / "context_before.json" + snapshot_path = output_dir / "source_snapshot.tar.gz" + if ( + (context_path.exists() or snapshot_path.exists()) + and not args.overwrite_output + ): + raise FileExistsError( + f"Context or source snapshot already exists in: {output_dir}" + ) + context = capture_context( + output_dir, + args.device, + use_runner_start_git=True, + ) + if context["environment"]["gpu"] is None: + raise ValueError( + "CUDA device evidence is missing from the pre-run context" + ) + write_source_snapshot( + list(context["source_file_sha256"]), + snapshot_path, + ) + context["source_snapshot"] = validate_source_snapshot( + snapshot_path, + context["source_file_sha256"], + ) + write_atomic( + context_path, + json.dumps( + context, + indent=2, + sort_keys=True, + allow_nan=False, + ) + + "\n", + ) + print(f"Wrote pre-run context snapshot to {context_path}") + return + + result_path = output_dir / "benchmark_report_results.json" + markdown_path = output_dir / "benchmark_report.md" + figure_path = output_dir / "benchmark_results.png" + delta_dpd_figure_path = output_dir / "benchmark_delta_dpd_results.png" + if ( + ( + result_path.exists() + or markdown_path.exists() + or figure_path.exists() + or delta_dpd_figure_path.exists() + ) + and not args.overwrite_output + ): + raise FileExistsError( + "Collector output already exists; use a new run directory or " + "--overwrite-output." + ) + + initial_context, completion_context = load_and_validate_context( + output_dir, args.device + ) + manifest = collect_manifest( + output_dir, + args.device, + initial_context, + completion_context, + ) + write_results_figure(manifest, figure_path) + write_delta_dpd_figure(manifest, delta_dpd_figure_path) + write_atomic( + markdown_path, + render_markdown(manifest), + ) + bind_derived_artifacts( + manifest, + report_path=markdown_path, + figure_path=figure_path, + delta_dpd_figure_path=delta_dpd_figure_path, + ) + write_atomic( + result_path, + json.dumps( + manifest, + indent=2, + sort_keys=True, + allow_nan=False, + ) + + "\n", + ) + print(f"Validated benchmark evidence in {output_dir}") + print(f"Wrote {output_dir / 'benchmark_report_results.json'}") + print(f"Wrote {output_dir / 'benchmark_report.md'}") + print(f"Wrote {output_dir / 'benchmark_results.png'}") + print(f"Wrote {output_dir / 'benchmark_delta_dpd_results.png'}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/reproduce_benchmark_report.sh b/benchmark/reproduce_benchmark_report.sh new file mode 100755 index 0000000..9050389 --- /dev/null +++ b/benchmark/reproduce_benchmark_report.sh @@ -0,0 +1,703 @@ +#!/usr/bin/env bash +# +# Reproduce benchmark/benchmark_report.md from the two bundled RF datasets. +# +# The benchmark matrix is fixed before execution: +# PA (~2,700 real parameters): MP, GMP, GRU, TRes-GRU, +# TRes-DeltaGRU (THX=THH=0) +# DPD (~1,000 real parameters): MP/ILA, GMP/ILA, GRU, TRes-GRU +# PA-surrogate sensitivity: TRes-DeltaGRU DPD (THX=THH=0) trained +# independently through TRes-GRU and +# TRes-DeltaGRU PA models +# +# All neural jobs use the same 300-epoch recipe. The primary DPD comparison +# uses the dataset's validation-selected TRes-GRU PA; the sensitivity runs use +# each of the two named PA surrogates. +# +# Default invocation: +# +# bash benchmark/reproduce_benchmark_report.sh +# +# Evidence is written to an immutable timestamped directory under +# benchmark/results/reproduced/. Architecture-addressed save/log artifacts are +# never overwritten unless --overwrite is supplied; in that case, the exact +# colliding files are archived into the run directory first. + +set -Eeuo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd) + +dry_run=0 +overwrite=0 +device=${DEVICE:-0} +output_dir="" + +usage() { + echo "Usage: $0 [--dry-run] [--overwrite] [--device INDEX] [--output-dir PATH]" + echo + echo " --dry-run Record and print every job without executing it." + echo " --overwrite Archive, then replace exact colliding save/log files." + echo " --device INDEX CUDA device index (default: ${device})." + echo " --output-dir PATH Run-specific evidence directory." +} + +while (($#)); do + case "$1" in + --dry-run) + dry_run=1 + shift + ;; + --overwrite) + overwrite=1 + shift + ;; + --device) + if (($# < 2)); then + echo "[ERROR] --device requires an index." >&2 + exit 2 + fi + device=$2 + shift 2 + ;; + --output-dir) + if (($# < 2)); then + echo "[ERROR] --output-dir requires a path." >&2 + exit 2 + fi + output_dir=$2 + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "[ERROR] Unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +if [[ -n "${OPENDPD_PYTHON:-}" ]]; then + python_bin=${OPENDPD_PYTHON} +elif [[ -n "${PYTHON:-}" ]]; then + python_bin=${PYTHON} +elif [[ -x "${REPO_ROOT}/.venv/bin/python" ]]; then + python_bin=${REPO_ROOT}/.venv/bin/python +elif command -v python3 >/dev/null 2>&1; then + python_bin=$(command -v python3) +else + echo "[ERROR] Python was not found. Set OPENDPD_PYTHON." >&2 + exit 1 +fi + +if [[ "${python_bin}" == */* ]]; then + if [[ "${python_bin}" != /* ]]; then + python_dir=$(cd "$(dirname "${python_bin}")" && pwd) + python_bin="${python_dir}/$(basename "${python_bin}")" + fi + if [[ ! -x "${python_bin}" ]]; then + echo "[ERROR] Python interpreter is not executable: ${python_bin}" >&2 + exit 1 + fi +else + resolved_python=$(type -P -- "${python_bin}" || true) + if [[ -z "${resolved_python}" ]]; then + echo "[ERROR] Python interpreter was not found on PATH: ${python_bin}" >&2 + exit 1 + fi + python_bin=${resolved_python} +fi + +if [[ ! "${device}" =~ ^[0-9]+$ ]]; then + echo "[ERROR] --device must be a non-negative integer: ${device}" >&2 + exit 2 +fi + +if [[ -z "${output_dir}" ]]; then + run_id=${RUN_ID:-$(date -u +%Y%m%dT%H%M%SZ)} + output_dir="${REPO_ROOT}/benchmark/results/reproduced/${run_id}" +elif [[ "${output_dir}" != /* ]]; then + output_dir="${REPO_ROOT}/${output_dir}" +fi + +benchmark_lock_fd="" +if ((dry_run == 0)); then + if ! command -v flock >/dev/null 2>&1; then + echo "[ERROR] flock is required to serialize benchmark reproductions." >&2 + exit 1 + fi + if ! benchmark_lock_path=$( + git -C "${REPO_ROOT}" rev-parse --git-path opendpd-benchmark-report.lock + ); then + echo "[ERROR] Could not resolve the repository benchmark lock path." >&2 + exit 1 + fi + if [[ "${benchmark_lock_path}" != /* ]]; then + benchmark_lock_path="${REPO_ROOT}/${benchmark_lock_path}" + fi + benchmark_lock_dir=$(cd "$(dirname "${benchmark_lock_path}")" && pwd) + benchmark_lock_path="${benchmark_lock_dir}/$(basename "${benchmark_lock_path}")" + exec {benchmark_lock_fd}> "${benchmark_lock_path}" + if ! flock -n "${benchmark_lock_fd}"; then + echo "[ERROR] Another benchmark reproduction is already running for this repository." >&2 + echo "Lock: ${benchmark_lock_path}" >&2 + exit 1 + fi +fi + +if [[ -e "${output_dir}" ]]; then + echo "[ERROR] Output path already exists: ${output_dir}" >&2 + echo "Choose a new --output-dir or RUN_ID; run directories are immutable." >&2 + exit 1 +fi + +for dataset in APA_200MHz DPA_160MHz; do + if [[ ! -f "${REPO_ROOT}/datasets/${dataset}/spec.json" ]]; then + echo "[ERROR] Required dataset is missing: datasets/${dataset}" >&2 + exit 1 + fi +done + +git_commit_before=$(git -C "${REPO_ROOT}" rev-parse HEAD) +git_branch_before=$(git -C "${REPO_ROOT}" branch --show-current) +git_status_before=$(git -C "${REPO_ROOT}" status --short --untracked-files=all) + +mkdir -p "$(dirname "${output_dir}")" +mkdir "${output_dir}" +mkdir "${output_dir}/logs" +commands_log="${output_dir}/commands.log" +jobs_log="${output_dir}/jobs.tsv" +pa_bindings_log="${output_dir}/pa_checkpoint_bindings.tsv" +: > "${commands_log}" +printf 'label\tstarted_at_utc\tfinished_at_utc\tduration_seconds\texit_status\n' > "${jobs_log}" +printf 'dataset\tdpd_model_key\tpa_model_key\tcheckpoint\tsha256\n' \ + > "${pa_bindings_log}" +printf '%s\n' "${git_commit_before}" > "${output_dir}/git_commit_before.txt" +printf '%s\n' "${git_branch_before}" > "${output_dir}/git_branch_before.txt" +printf '%s\n' "${git_status_before}" > "${output_dir}/git_status_before.txt" + +cd "${REPO_ROOT}" + +"${python_bin}" - "${output_dir}/recipe.json" <<'PY' +import json +from pathlib import Path +import sys + +recipe = { + "schema_version": 6, + "datasets": ["APA_200MHz", "DPA_160MHz"], + "parameter_count_convention": ( + "Neural parameters are real trainable scalars; each complex polynomial " + "coefficient counts as two real degrees of freedom." + ), + "neural": { + "optimizer": "AdamW", + "optimizer_weight_decay": 0.01, + "optimizer_betas": [0.9, 0.999], + "optimizer_eps": 1e-8, + "loss": "MSE", + "epochs": 300, + "batch_size": 64, + "evaluation_batch_size": 64, + "initial_learning_rate": 5e-3, + "scheduler": "ReduceLROnPlateau", + "scheduler_factor": 0.5, + "scheduler_patience": 5, + "scheduler_threshold": 1e-4, + "scheduler_threshold_mode": "rel", + "scheduler_cooldown": 0, + "scheduler_eps": 1e-8, + "minimum_learning_rate": 5e-5, + "frame_length": 200, + "frame_stride": 1, + "gradient_clip": 200, + "seed": 0, + "reproducibility": "soft", + "pa_scheduler_and_selection_metric": "validation NMSE", + "dpd_scheduler_and_selection_metric": "validation ACLR_AVG", + }, + "pa_models": { + "mp": {"K": 9, "Q": 150, "real_parameters": 2700}, + "gmp": { + "Ka": 5, "La": 30, + "Kb": 4, "Lb": 30, "Mb": 5, + "Kc": 4, "Lc": 30, "Mc": 5, + "real_parameters": 2700, + }, + "gru": {"hidden_size": 28, "real_parameters": 2746}, + "tres_gru": {"hidden_size": 27, "real_parameters": 2751}, + "tres_deltagru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + }, + }, + "dpd_models": { + "mp": {"K": 5, "Q": 100, "real_parameters": 1000, "method": "ILA"}, + "gmp": { + "Ka": 5, "La": 20, + "Kb": 4, "Lb": 20, "Mb": 3, + "Kc": 4, "Lc": 20, "Mc": 2, + "real_parameters": 1000, + "method": "ILA", + }, + "gru": {"hidden_size": 16, "real_parameters": 994}, + "tres_gru": {"hidden_size": 15, "real_parameters": 999}, + }, + "dpd_tres_deltagru_pa_surrogate_comparison": { + "model": { + "hidden_size": 15, + "real_parameters": 999, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + }, + "pa_surrogates": { + "tres_gru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": None, + }, + "tres_deltagru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + }, + }, + }, + "dpd_pa_surrogate": { + "backbone": "tres_gru", + "hidden_size": 27, + "real_parameters": 2751, + }, + "polynomial_solvers": { + "default": { + "implementation": "torch.linalg.lstsq", + "mode": "gels", + "driver": "gels", + "dtype": "torch.complex64", + "column_scaling": "l2", + "regularization": None, + }, + "gmp_pa": { + "implementation": "torch.linalg.svd", + "mode": "truncated_svd", + "driver": "gesvdj", + "dtype": "torch.complex64", + "column_scaling": "l2", + "regularization": { + "type": "truncated_svd", + "relative_cutoff": 1e-4, + }, + }, + "segment_boundary_policy": "zero delay state at every nperseg boundary", + }, +} + +path = Path(sys.argv[1]) +path.write_text(json.dumps(recipe, indent=2, sort_keys=True) + "\n") +PY + +record_command() { + local label=$1 + shift + printf '# %s\n' "${label}" >> "${commands_log}" + printf '%q ' "$@" >> "${commands_log}" + printf '\n\n' >> "${commands_log}" +} + +print_command() { + local label=$1 + shift + printf '\n[%s]\n' "${label}" + printf '+ ' + printf '%q ' "$@" + printf '\n' +} + +run_step() { + local label=$1 + shift + record_command "${label}" "$@" + print_command "${label}" "$@" + if ((dry_run)); then + return + fi + + local started_at finished_at started_epoch finished_epoch duration status + started_at=$(date -u +%Y-%m-%dT%H:%M:%SZ) + started_epoch=$(date +%s) + set +e + "$@" 2>&1 | tee "${output_dir}/logs/${label}.log" + status=${PIPESTATUS[0]} + set -e + finished_epoch=$(date +%s) + finished_at=$(date -u +%Y-%m-%dT%H:%M:%SZ) + duration=$((finished_epoch - started_epoch)) + printf '%s\t%s\t%s\t%s\t%s\n' \ + "${label}" "${started_at}" "${finished_at}" "${duration}" "${status}" \ + >> "${jobs_log}" + if ((status != 0)); then + echo "[ERROR] Job failed (${status}): ${label}" >&2 + exit "${status}" + fi +} + +copy_with_parents() { + local destination=$1 + shift + local path + for path in "$@"; do + [[ -f "${path}" ]] || continue + mkdir -p "${destination}/$(dirname "${path}")" + cp -p "${path}" "${destination}/${path}" + done +} + +sha256_path() { + "${python_bin}" - "$1" <<'PY' +import hashlib +from pathlib import Path +import sys + +path = Path(sys.argv[1]) +if not path.is_file(): + raise SystemExit(f"[ERROR] Required checkpoint is missing: {path}") +digest = hashlib.sha256() +with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) +print(digest.hexdigest()) +PY +} + +record_pa_binding() { + local dataset=$1 + local dpd_model_key=$2 + local pa_model_key=$3 + local checkpoint=$4 + if ((dry_run)); then + return + fi + local digest + digest=$(sha256_path "${checkpoint}") + printf '%s\t%s\t%s\t%s\t%s\n' \ + "${dataset}" "${dpd_model_key}" "${pa_model_key}" \ + "${checkpoint}" "${digest}" \ + >> "${pa_bindings_log}" +} + +# Only these exact architecture-addressed files can collide with this recipe. +target_artifacts=() +add_training_artifacts() { + local dataset=$1 + local step=$2 + local parent=$3 + local model_id=$4 + local save_base="save/${dataset}/${step}" + local log_base="log/${dataset}/${step}" + if [[ -n "${parent}" ]]; then + save_base="${save_base}/${parent}" + log_base="${log_base}/${parent}" + fi + target_artifacts+=( + "${save_base}/${model_id}.pt" + "${log_base}/best/${model_id}.csv" + "${log_base}/history/${model_id}.csv" + ) +} + +for dataset in APA_200MHz DPA_160MHz; do + add_training_artifacts "${dataset}" train_pa "" \ + PA_S_0_M_GRU_H_28_F_200_P_2746 + add_training_artifacts "${dataset}" train_pa "" \ + PA_S_0_M_TRES_GRU_H_27_F_200_P_2751 + add_training_artifacts "${dataset}" train_pa "" \ + PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751 + + dpd_parent=PA_S_0_M_TRES_GRU_H_27_F_200 + add_training_artifacts "${dataset}" train_dpd "${dpd_parent}" \ + DPD_S_0_M_GRU_H_16_F_200_P_994 + add_training_artifacts "${dataset}" train_dpd "${dpd_parent}" \ + DPD_S_0_M_TRES_GRU_H_15_F_200_P_999 + add_training_artifacts "${dataset}" train_dpd "${dpd_parent}" \ + DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000 + + delta_dpd_parent=PA_S_0_M_TRES_DELTAGRU_H_27_F_200 + add_training_artifacts "${dataset}" train_dpd "${delta_dpd_parent}" \ + DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000 +done + +existing_artifacts=() +for path in "${target_artifacts[@]}"; do + if [[ -L "${path}" ]]; then + echo "[ERROR] Refusing benchmark artifact symlink: ${path}" >&2 + exit 1 + elif [[ -e "${path}" && ! -f "${path}" ]]; then + echo "[ERROR] Refusing non-regular benchmark artifact: ${path}" >&2 + exit 1 + elif [[ -f "${path}" ]]; then + existing_artifacts+=("${path}") + fi +done + +if ((dry_run == 0 && ${#existing_artifacts[@]} > 0 && overwrite == 0)); then + echo "[ERROR] Matching benchmark artifacts already exist and would be overwritten:" >&2 + printf ' %s\n' "${existing_artifacts[@]}" >&2 + echo "Re-run with --overwrite to archive these exact files before training." >&2 + exit 1 +fi + +if ((dry_run == 0 && ${#existing_artifacts[@]} > 0)); then + copy_with_parents "${output_dir}/preexisting" "${existing_artifacts[@]}" + echo "[INFO] Archived ${#existing_artifacts[@]} pre-existing artifact(s)." +fi + +if ((dry_run == 0)); then + "${python_bin}" - "${device}" <<'PY' +import sys + +try: + import numpy + import pandas + import scipy + import torch +except ImportError as exc: + raise SystemExit(f"[ERROR] Missing benchmark dependency: {exc}") from exc + +from models import CoreModel +from utils.util import count_net_params + +device_index = int(sys.argv[1]) +if not torch.cuda.is_available(): + raise SystemExit("[ERROR] CUDA is required for the reference benchmark.") +if device_index >= torch.cuda.device_count(): + raise SystemExit( + f"[ERROR] CUDA device {device_index} is unavailable; " + f"detected {torch.cuda.device_count()} device(s)." + ) +device = torch.device(f"cuda:{device_index}") +print(f"[INFO] Python {sys.version.split()[0]}") +print(f"[INFO] PyTorch {torch.__version__}; CUDA {torch.version.cuda}") +print(f"[INFO] CUDA device {device_index}: {torch.cuda.get_device_name(device)}") + +expected = { + ("gru", 28): 2746, + ("tres_gru", 27): 2751, + ("tres_deltagru", 27): 2751, + ("gru", 16): 994, + ("tres_gru", 15): 999, + ("tres_deltagru", 15): 999, +} +for (backbone, hidden_size), expected_parameters in expected.items(): + network = CoreModel( + input_size=2, + hidden_size=hidden_size, + num_layers=1, + backbone_type=backbone, + thx=0.0, + thh=0.0, + ).to(device) + actual_parameters = count_net_params(network) + if actual_parameters != expected_parameters: + raise SystemExit( + f"[ERROR] {backbone}-H{hidden_size} has {actual_parameters} " + f"parameters, expected {expected_parameters}." + ) + sample = torch.randn(2, 200, 2, device=device, requires_grad=True) + output = network(sample) + output.square().mean().backward() + if not torch.isfinite(output).all() or sample.grad is None: + raise SystemExit(f"[ERROR] Non-finite smoke test for {backbone}-H{hidden_size}.") + del network, sample, output + +matrix = torch.randn(512, 16, dtype=torch.complex64, device=device) +target = torch.randn(512, 1, dtype=torch.complex64, device=device) +solution = torch.linalg.lstsq(matrix, target, driver="gels").solution +if not torch.isfinite(solution).all(): + raise SystemExit("[ERROR] CUDA complex least-squares smoke test failed.") +torch.cuda.synchronize(device) +print("[INFO] Neural and complex least-squares CUDA smoke tests passed.") +PY +fi + +neural_common=( + --PA_num_layers 1 + --DPD_num_layers 1 + --frame_length 200 + --frame_stride 1 + --n_epochs 300 + --opt_type adamw + --lr 5e-3 + --lr_schedule 1 + --lr_end 5e-5 + --decay_factor 0.5 + --patience 5 + --batch_size 64 + --batch_size_eval 64 + --loss_type l2 + --grad_clip_val 200 + --thx 0 + --thh 0 + --log_precision 8 + --seed 0 + --re_level soft + --eval_val 1 + --eval_test 1 + --accelerator cuda + --devices "${device}" +) + +run_step snapshot_context \ + "${python_bin}" -m benchmark.collect_benchmark_report \ + --output-dir "${output_dir}" --device "${device}" --snapshot-context + +run_neural_pa() { + local dataset=$1 + local slug=$2 + local backbone=$3 + local hidden_size=$4 + run_step "train_pa_${slug}_${backbone}" \ + "${python_bin}" main.py \ + --dataset_name "${dataset}" \ + --step train_pa \ + --PA_backbone "${backbone}" \ + --PA_hidden_size "${hidden_size}" \ + "${neural_common[@]}" +} + +run_neural_dpd() { + local dataset=$1 + local slug=$2 + local backbone=$3 + local hidden_size=$4 + local pa_model_key=$5 + local pa_backbone=$6 + local pa_hidden_size=$7 + local pa_model_id=$8 + local pa_checkpoint="save/${dataset}/train_pa/${pa_model_id}.pt" + record_pa_binding \ + "${dataset}" "${backbone}" "${pa_model_key}" "${pa_checkpoint}" + run_step "train_dpd_${slug}_${backbone}_via_${pa_model_key}_pa" \ + "${python_bin}" main.py \ + --dataset_name "${dataset}" \ + --step train_dpd \ + --PA_backbone "${pa_backbone}" \ + --PA_hidden_size "${pa_hidden_size}" \ + --DPD_backbone "${backbone}" \ + --DPD_hidden_size "${hidden_size}" \ + "${neural_common[@]}" +} + +run_polynomial_pa() { + local dataset=$1 + local slug=$2 + local polynomial=$3 + shift 3 + local solver_arguments=(--solver-mode gels) + if [[ "${polynomial}" == "gmp" ]]; then + solver_arguments=(--solver-mode truncated_svd --svd-rcond 1e-4) + fi + run_step "pa_model_${slug}_${polynomial}" \ + "${python_bin}" -m benchmark.benchmark_volterra \ + --task pa_modeling \ + --dataset-name "${dataset}" \ + --model "${polynomial}" \ + --solver-device "cuda:${device}" \ + --solver-dtype complex64 \ + "${solver_arguments[@]}" \ + --device "${device}" \ + --json-out "${output_dir}/benchmark_report_${slug}_pa_${polynomial}.json" \ + "$@" +} + +run_polynomial_dpd() { + local dataset=$1 + local slug=$2 + local polynomial=$3 + shift 3 + run_step "dpd_ila_${slug}_${polynomial}" \ + "${python_bin}" -m benchmark.benchmark_volterra \ + --task dpd_ila \ + --dataset-name "${dataset}" \ + --model "${polynomial}" \ + --solver-device "cuda:${device}" \ + --solver-dtype complex64 \ + --solver-mode gels \ + --device "${device}" \ + --pa-backbone tres_gru \ + --pa-hidden-size 27 \ + --pa-num-layers 1 \ + --pa-checkpoint \ + "save/${dataset}/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt" \ + --json-out "${output_dir}/benchmark_report_${slug}_dpd_${polynomial}.json" \ + "$@" +} + +for dataset_spec in "APA_200MHz apa" "DPA_160MHz dpa"; do + read -r dataset slug <<< "${dataset_spec}" + run_neural_pa "${dataset}" "${slug}" gru 28 + run_neural_pa "${dataset}" "${slug}" tres_gru 27 + run_neural_pa "${dataset}" "${slug}" tres_deltagru 27 +done + +for dataset_spec in "APA_200MHz apa" "DPA_160MHz dpa"; do + read -r dataset slug <<< "${dataset_spec}" + run_polynomial_pa "${dataset}" "${slug}" mp \ + --K 9 --Q 150 + run_polynomial_pa "${dataset}" "${slug}" gmp \ + --Ka 5 --La 30 --Kb 4 --Lb 30 --Mb 5 \ + --Kc 4 --Lc 30 --Mc 5 +done + +for dataset_spec in "APA_200MHz apa" "DPA_160MHz dpa"; do + read -r dataset slug <<< "${dataset_spec}" + run_neural_dpd \ + "${dataset}" "${slug}" gru 16 tres_gru tres_gru 27 \ + PA_S_0_M_TRES_GRU_H_27_F_200_P_2751 + run_neural_dpd \ + "${dataset}" "${slug}" tres_gru 15 tres_gru tres_gru 27 \ + PA_S_0_M_TRES_GRU_H_27_F_200_P_2751 + run_neural_dpd \ + "${dataset}" "${slug}" tres_deltagru 15 tres_gru tres_gru 27 \ + PA_S_0_M_TRES_GRU_H_27_F_200_P_2751 + run_neural_dpd \ + "${dataset}" "${slug}" tres_deltagru 15 \ + tres_deltagru tres_deltagru 27 \ + PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751 +done + +for dataset_spec in "APA_200MHz apa" "DPA_160MHz dpa"; do + read -r dataset slug <<< "${dataset_spec}" + run_polynomial_dpd "${dataset}" "${slug}" mp \ + --K 5 --Q 100 + run_polynomial_dpd "${dataset}" "${slug}" gmp \ + --Ka 5 --La 20 --Kb 4 --Lb 20 --Mb 3 \ + --Kc 4 --Lc 20 --Mc 2 +done + +if ((dry_run)); then + record_command collect_report \ + "${python_bin}" -m benchmark.collect_benchmark_report \ + --output-dir "${output_dir}" --device "${device}" + print_command collect_report \ + "${python_bin}" -m benchmark.collect_benchmark_report \ + --output-dir "${output_dir}" --device "${device}" + echo + echo "[DRY RUN] Commands were written to ${commands_log}" + exit 0 +fi + +copy_with_parents "${output_dir}/artifacts" "${target_artifacts[@]}" + +run_step collect_report \ + "${python_bin}" -m benchmark.collect_benchmark_report \ + --output-dir "${output_dir}" --device "${device}" + +echo +echo "[DONE] Reproduced benchmark evidence: ${output_dir}" +echo "[DONE] Markdown report: ${output_dir}/benchmark_report.md" +echo "[DONE] Machine-readable results: ${output_dir}/benchmark_report_results.json" +echo "[DONE] Results visualization: ${output_dir}/benchmark_results.png" +echo "[DONE] Delta-DPD sensitivity visualization: ${output_dir}/benchmark_delta_dpd_results.png" diff --git a/benchmark/results/benchmark_report_results.json b/benchmark/results/benchmark_report_results.json new file mode 100644 index 0000000..36f970f --- /dev/null +++ b/benchmark/results/benchmark_report_results.json @@ -0,0 +1,1904 @@ +{ + "artifact_kind": "opendpd_benchmark_reproduction_run", + "canonical_report": "benchmark/benchmark_report.md", + "completion_context": { + "captured_at_utc": "2026-07-26T20:45:34.695466+00:00", + "git_branch": "benchmark-fix", + "git_commit": "3df35e081e6e41463fa46f21778c72a748823274", + "validated_unchanged": true + }, + "datasets": { + "APA_200MHz": { + "dpd_models": { + "gmp": { + "artifact": { + "path": "benchmark_report_apa_dpd_gmp.json", + "sha256": "62912b11d4b370e2a0ae68f60bb03b53351e8c264bc0ddad149101beceb953d1", + "size_bytes": 2972 + }, + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 20, + "Lb": 20, + "Lc": 20, + "Mb": 3, + "Mc": 2 + }, + "coefficient_l2_norm": 46.2047004699707, + "column_scale_max": 89.97786712646484, + "column_scale_min": 9.357341766357422, + "column_scaling": "l2", + "complex_coefficients": 500, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "indirect_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -43.5931687089929, + "aclr_left_db": -44.06568580030243, + "aclr_right_db": -43.120651617683365, + "evm_db": -46.352304380006316, + "nmse_db": -38.52589797973633 + }, + "validation": { + "aclr_avg_db": -42.972747788860104, + "aclr_left_db": -43.55596870580729, + "aclr_right_db": -42.38952687191292, + "evm_db": -46.27533200387405, + "nmse_db": -38.31670379638672 + } + }, + "model": { + "backbone": "gmp", + "display_name": "GMP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 1000 + }, + "nperseg": 19662, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb" + }, + "pa_evaluation_device": "cuda:0", + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 1000, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 983040000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 1.0, + "task": "dpd_ila", + "training_relative_residual": 0.010865671560168266 + }, + "gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_GRU_H_16_F_200_P_994.csv", + "sha256": "a6d38797ef37bbe21ce76ee9b5068b8b158172c1b3ba10729e9095278812d2e6", + "size_bytes": 421, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_GRU_H_16_F_200_P_994.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_GRU_H_16_F_200_P_994.pt", + "sha256": "e20f7c4c97ed09ea991b479812db21bef9ebdc622d4a1ad1bb2418988602d627", + "size_bytes": 6832, + "source_path": "save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_GRU_H_16_F_200_P_994.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_GRU_H_16_F_200_P_994.csv", + "sha256": "f535a28e527eb0594ac0f01450ed8cbfdb685c72f19aee8705005b533e030a3e", + "size_bytes": 62248, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_GRU_H_16_F_200_P_994.csv" + } + }, + "completed_epochs": 300, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -51.0114043, + "aclr_left_db": -51.0559093, + "aclr_right_db": -50.9668993, + "evm_db": -47.42746888, + "loss": 1.52e-06, + "nmse_db": -45.13145 + }, + "validation": { + "aclr_avg_db": -51.32400487, + "aclr_left_db": -51.38568411, + "aclr_right_db": -51.26232562, + "evm_db": -47.68851727, + "loss": 1.5e-06, + "nmse_db": -45.17349 + } + }, + "model": { + "backbone": "gru", + "display_name": "GRU-H16", + "hidden_size": 16, + "model_id": "DPD_S_0_M_GRU_H_16_F_200_P_994", + "parameters": 994 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "APA_200MHz", + "dpd_model_key": "gru", + "pa_model_key": "tres_gru", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb" + }, + "selected_epoch_zero_based": 299, + "selected_learning_rate": 0.00015625, + "task": "dpd_training", + "train_loss": 2.17e-06 + }, + "mp": { + "artifact": { + "path": "benchmark_report_apa_dpd_mp.json", + "sha256": "8d24721c1e3783f2247b4906e9b50ecb2feefe0040e0c93cf5ea9c6562ee3640", + "size_bytes": 2893 + }, + "basis_configuration": { + "K": 5, + "Q": 100 + }, + "coefficient_l2_norm": 4.383342266082764, + "column_scale_max": 89.97786712646484, + "column_scale_min": 13.769466400146484, + "column_scaling": "l2", + "complex_coefficients": 500, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "indirect_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -45.185983122836305, + "aclr_left_db": -46.0739531814713, + "aclr_right_db": -44.2980130642013, + "evm_db": -48.15347693703668, + "nmse_db": -42.189632415771484 + }, + "validation": { + "aclr_avg_db": -44.922360391352086, + "aclr_left_db": -45.558971707171324, + "aclr_right_db": -44.28574907553285, + "evm_db": -48.18950883574672, + "nmse_db": -42.11494827270508 + } + }, + "model": { + "backbone": "mp", + "display_name": "MP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 1000 + }, + "nperseg": 19662, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb" + }, + "pa_evaluation_device": "cuda:0", + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 1000, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 983040000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 1.0, + "task": "dpd_ila", + "training_relative_residual": 0.011005041189491749 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv", + "sha256": "60137327bb949de2422c86587071a005d68ea78d94887801156a24c995483880", + "size_bytes": 427, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.pt", + "sha256": "b26552ffbafe8edcdfb6df66ec50f179b81644977f6c2f331f471d3f76b2d9fc", + "size_bytes": 7541, + "source_path": "save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv", + "sha256": "11afe746632723165fbad0abc4fa4a854c5aa41587aa0d749939feac589de9e9", + "size_bytes": 63720, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv" + } + }, + "completed_epochs": 300, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -53.48793859, + "aclr_left_db": -53.46930447, + "aclr_right_db": -53.50657272, + "evm_db": -45.09727783, + "loss": 1.84e-06, + "nmse_db": -44.285923 + }, + "validation": { + "aclr_avg_db": -53.47358387, + "aclr_left_db": -53.20646222, + "aclr_right_db": -53.74070552, + "evm_db": -44.86538699, + "loss": 1.95e-06, + "nmse_db": -44.03681 + } + }, + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H15", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_GRU_H_15_F_200_P_999", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "APA_200MHz", + "dpd_model_key": "tres_gru", + "pa_model_key": "tres_gru", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb" + }, + "selected_epoch_zero_based": 296, + "selected_learning_rate": 5e-05, + "task": "dpd_training", + "train_loss": 3.05e-06 + } + }, + "dpd_tres_deltagru_by_pa_surrogate": { + "tres_deltagru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "deb1fd39e3fde771614e01641ae9af14f96be47df3343611a8a941f5fe53e478", + "size_bytes": 462, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt", + "sha256": "3699df3ad1fee950281b945e3f43f40015ab6f79380113b35ba981a1cff6bf5c", + "size_bytes": 8160, + "source_path": "save/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "61e1ea3e9d31edd7d8cd63be2994f52b7b636d359bc8907d33138ac8a1c5aa16", + "size_bytes": 71889, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -54.40097617, + "aclr_left_db": -54.56885076, + "aclr_right_db": -54.23310159, + "evm_db": -49.69994504, + "loss": 7.9e-07, + "nmse_db": -47.970295 + }, + "validation": { + "aclr_avg_db": -54.48738132, + "aclr_left_db": -55.10222104, + "aclr_right_db": -53.87254159, + "evm_db": -49.82003181, + "loss": 8.1e-07, + "nmse_db": -47.85823 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "dataset": "APA_200MHz", + "dpd_model_key": "tres_deltagru", + "pa_model_key": "tres_deltagru", + "sha256": "b9451d33c6c4433e5c5d5a1326ab80299fe4ac51d07329e00dab183eece22650" + }, + "pa_surrogate": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "b9451d33c6c4433e5c5d5a1326ab80299fe4ac51d07329e00dab183eece22650", + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H27 (THX=THH=0)", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "model_key": "tres_deltagru", + "selection_metric": "validation NMSE" + }, + "selected_epoch_zero_based": 296, + "selected_learning_rate": 5e-05, + "task": "dpd_training", + "train_loss": 1.51e-06 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "88afdf9d220f1445798caf7722627d4396e52849a04fd961c03a75bc3de0ca79", + "size_bytes": 462, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt", + "sha256": "8fc444b84bf85c7935f9d842b8753a8cdb9030672d0d5c61a5a1d8fb45d0f411", + "size_bytes": 8160, + "source_path": "save/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "fa350b8ac3620ed1abba12adffe797ca3c5c81fe3d10dbfd7488fe139369970a", + "size_bytes": 71864, + "source_path": "log/APA_200MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -53.64353511, + "aclr_left_db": -53.87717378, + "aclr_right_db": -53.40989644, + "evm_db": -48.26290132, + "loss": 1.01e-06, + "nmse_db": -46.91838 + }, + "validation": { + "aclr_avg_db": -53.2897953, + "aclr_left_db": -53.47258954, + "aclr_right_db": -53.10700107, + "evm_db": -48.34864182, + "loss": 1.02e-06, + "nmse_db": -46.847137 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "APA_200MHz", + "dpd_model_key": "tres_deltagru", + "pa_model_key": "tres_gru", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb" + }, + "pa_surrogate": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb", + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H27", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "model_key": "tres_gru", + "selection_metric": "validation NMSE" + }, + "selected_epoch_zero_based": 264, + "selected_learning_rate": 5e-05, + "task": "dpd_training", + "train_loss": 2.13e-06 + } + }, + "pa_models": { + "gmp": { + "artifact": { + "path": "benchmark_report_apa_pa_gmp.json", + "sha256": "e9e22d371ce762aa35667760df5cfb782747f0b6b6f9c6391e7f807dce421943", + "size_bytes": 2803 + }, + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 30, + "Lb": 30, + "Lc": 30, + "Mb": 5, + "Mc": 5 + }, + "coefficient_l2_norm": 7.8917012214660645, + "column_scale_max": 76.96764373779297, + "column_scale_min": 3.718569040298462, + "column_scaling": "l2", + "complex_coefficients": 1350, + "condition_number": 4101974003.5174985, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": false, + "least_squares_rank": 650, + "method": "direct_least_squares", + "metrics": { + "test": { + "aclr_avg_db": -27.70622397315019, + "aclr_left_db": -27.783035363162686, + "aclr_right_db": -27.629412583137697, + "evm_db": -42.797116668332876, + "nmse_db": -38.66061984321832 + }, + "validation": { + "aclr_avg_db": -27.409084690382358, + "aclr_left_db": -27.40526404384027, + "aclr_right_db": -27.412905336924446, + "evm_db": -43.053951759619196, + "nmse_db": -38.7019771115391 + } + }, + "model": { + "backbone": "gmp", + "display_name": "GMP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 2700 + }, + "nperseg": 19662, + "pa_checkpoint_binding": null, + "pa_evaluation_device": null, + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 2700, + "regularization": "truncated_svd", + "retained_singular_value_min": 0.0013197693042457104, + "sample_rate_hz": 983040000.0, + "singular_value_cutoff": 0.0012831207275390626, + "singular_value_max": 12.831207275390625, + "singular_value_min": 3.1280567025504524e-09, + "solver": "torch.linalg.svd", + "solver_driver": "gesvdj", + "solver_mode": "truncated_svd", + "svd_rcond": 0.0001, + "target_gain": 1.0, + "task": "pa_modeling", + "training_relative_residual": 0.011340372264385223 + }, + "gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/best/PA_S_0_M_GRU_H_28_F_200_P_2746.csv", + "sha256": "7adb1c12e1a2e484c9426de6f31f70c9735e5b9e15efb686ec28be9bbbdc775d", + "size_bytes": 421, + "source_path": "log/APA_200MHz/train_pa/best/PA_S_0_M_GRU_H_28_F_200_P_2746.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_pa/PA_S_0_M_GRU_H_28_F_200_P_2746.pt", + "sha256": "71de3649529e335fde4d26d231728baf7e3d3b837d23a7e60768a62979312cb7", + "size_bytes": 13808, + "source_path": "save/APA_200MHz/train_pa/PA_S_0_M_GRU_H_28_F_200_P_2746.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/history/PA_S_0_M_GRU_H_28_F_200_P_2746.csv", + "sha256": "3ada71d8783b7ec1475ea68fe2bcd33bb4931ecc6452be639bd423c1f58a4c9d", + "size_bytes": 62253, + "source_path": "log/APA_200MHz/train_pa/history/PA_S_0_M_GRU_H_28_F_200_P_2746.csv" + } + }, + "completed_epochs": 300, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -27.67541149, + "aclr_left_db": -27.67831335, + "aclr_right_db": -27.67250962, + "evm_db": -43.57318158, + "loss": 8.64e-06, + "nmse_db": -38.93669 + }, + "validation": { + "aclr_avg_db": -27.38332431, + "aclr_left_db": -27.34550046, + "aclr_right_db": -27.42114816, + "evm_db": -43.63170309, + "loss": 8.78e-06, + "nmse_db": -38.85041 + } + }, + "model": { + "backbone": "gru", + "display_name": "GRU-H28", + "hidden_size": 28, + "model_id": "PA_S_0_M_GRU_H_28_F_200_P_2746", + "parameters": 2746 + }, + "selected_epoch_zero_based": 292, + "selected_learning_rate": 5e-05, + "task": "pa_modeling", + "train_loss": 1.486e-05 + }, + "mp": { + "artifact": { + "path": "benchmark_report_apa_pa_mp.json", + "sha256": "dc056f90d336dbaa5e6218cdf236d58133c191508e6ceb90303faccba63e1a6d", + "size_bytes": 2518 + }, + "basis_configuration": { + "K": 9, + "Q": 150 + }, + "coefficient_l2_norm": 1649.8330078125, + "column_scale_max": 76.96764373779297, + "column_scale_min": 3.0452980995178223, + "column_scaling": "l2", + "complex_coefficients": 1350, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "direct_least_squares", + "metrics": { + "test": { + "aclr_avg_db": -27.831701026650798, + "aclr_left_db": -27.886854192144813, + "aclr_right_db": -27.776547861156786, + "evm_db": -39.95707758634159, + "nmse_db": -36.963532286352084 + }, + "validation": { + "aclr_avg_db": -27.38812910025664, + "aclr_left_db": -27.405965138204177, + "aclr_right_db": -27.37029306230911, + "evm_db": -40.23019407942715, + "nmse_db": -37.040538398646866 + } + }, + "model": { + "backbone": "mp", + "display_name": "MP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 2700 + }, + "nperseg": 19662, + "pa_checkpoint_binding": null, + "pa_evaluation_device": null, + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 2700, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 983040000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 1.0, + "task": "pa_modeling", + "training_relative_residual": 0.013689383864402771 + }, + "tres_deltagru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/best/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv", + "sha256": "b71d9051ce4ee78ec8a223abf93cb2730c3f4bfa6e5a81d603bef1ec53f52410", + "size_bytes": 463, + "source_path": "log/APA_200MHz/train_pa/best/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "sha256": "b9451d33c6c4433e5c5d5a1326ab80299fe4ac51d07329e00dab183eece22650", + "size_bytes": 14596, + "source_path": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/history/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv", + "sha256": "3c938909f1f9c247499a7a07c401a7385680db95844788eca9a8e61cec294236", + "size_bytes": 71953, + "source_path": "log/APA_200MHz/train_pa/history/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -27.6760638, + "aclr_left_db": -27.68001247, + "aclr_right_db": -27.67211513, + "evm_db": -43.96833133, + "loss": 8.18e-06, + "nmse_db": -39.177734 + }, + "validation": { + "aclr_avg_db": -27.41638323, + "aclr_left_db": -27.3872018, + "aclr_right_db": -27.44556465, + "evm_db": -44.0254212, + "loss": 8.3e-06, + "nmse_db": -39.09286 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H27 (THX=THH=0)", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "selected_epoch_zero_based": 294, + "selected_learning_rate": 5e-05, + "task": "pa_modeling", + "train_loss": 1.417e-05 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/best/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv", + "sha256": "e0961441c4d10dd6292420c2032b7dff6083f24d7ba4976f7ef666351749fae5", + "size_bytes": 427, + "source_path": "log/APA_200MHz/train_pa/best/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb", + "size_bytes": 14773, + "source_path": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/APA_200MHz/train_pa/history/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv", + "sha256": "4a01e2a3108a95b63af42c5a9f1035cfd6992c53a651f60efb2b21a48fadd2be", + "size_bytes": 63772, + "source_path": "log/APA_200MHz/train_pa/history/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv" + } + }, + "completed_epochs": 300, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -27.69532568, + "aclr_left_db": -27.75026686, + "aclr_right_db": -27.64038451, + "evm_db": -43.92993785, + "loss": 8.27e-06, + "nmse_db": -39.12926 + }, + "validation": { + "aclr_avg_db": -27.40869864, + "aclr_left_db": -27.40999309, + "aclr_right_db": -27.40740419, + "evm_db": -43.91182816, + "loss": 8.39e-06, + "nmse_db": -39.044727 + } + }, + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H27", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "selected_epoch_zero_based": 288, + "selected_learning_rate": 5e-05, + "task": "pa_modeling", + "train_loss": 1.415e-05 + } + }, + "reference_pa": { + "checkpoint": "save/APA_200MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "fe08b5bdfa377af79bf1ea78fcfa0c96b10659bb050fe46ed2f3ef57dfcbc2eb", + "model_key": "tres_gru", + "selection_metric": "validation NMSE" + }, + "signal": { + "dataset_file_sha256": { + "datasets/APA_200MHz/test_input.csv": "5027d3d69391ed22ad79c410831bdfed47b25045088dda0756801cf591c947bf", + "datasets/APA_200MHz/test_output.csv": "991c1f97f38e57c7f3614dd3f23c411791e03d5eb5eae67492455cf9e94820f6", + "datasets/APA_200MHz/train_input.csv": "2ca703c9e7eb39839db1fb01f91081a86e18535e53751c587cc71d2d71e9c625", + "datasets/APA_200MHz/train_output.csv": "e760adf3908ed1be1e610c46f056e88bad6107a81cc8b01d91306727316b5930", + "datasets/APA_200MHz/val_input.csv": "39bb15c9bd92549d1653498c140caff5cb2f20edffd433eafa46b4a81c491981", + "datasets/APA_200MHz/val_output.csv": "02f67574444c7a8ba321cde1ea919c07fef1c99fb9d25678befc019c6b6645e2" + }, + "papr": { + "absolute_peak_db": 10.010544625499463, + "ccdf_db": 10.007981673494797, + "ccdf_probability": 1e-05, + "quantile_method": "numpy linear" + }, + "samples": 98304, + "target_gain": 1.0 + }, + "spec_sha256": "edaf447bc2d59c2cf10df4ad5bececf057bd8822a45864ec8959ddea1a0ebf9a" + }, + "DPA_160MHz": { + "dpd_models": { + "gmp": { + "artifact": { + "path": "benchmark_report_dpa_dpd_gmp.json", + "sha256": "56aef7fbcb816485d0d1339ed2a5bd068d609308a21d144203953ce2778328a2", + "size_bytes": 2998 + }, + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 20, + "Lb": 20, + "Lc": 20, + "Mb": 3, + "Mc": 2 + }, + "coefficient_l2_norm": 19.907581329345703, + "column_scale_max": 210.84939575195312, + "column_scale_min": 24.099092483520508, + "column_scaling": "l2", + "complex_coefficients": 500, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "indirect_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -53.112459466114444, + "aclr_left_db": -52.616658084847344, + "aclr_right_db": -53.608260847381544, + "evm_db": -50.188747977606816, + "nmse_db": -43.97211837768555 + }, + "validation": { + "aclr_avg_db": -52.78473824545991, + "aclr_left_db": -52.35719844738584, + "aclr_right_db": -53.212278043533985, + "evm_db": -50.12855174608197, + "nmse_db": -44.266239166259766 + } + }, + "model": { + "backbone": "gmp", + "display_name": "GMP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 1000 + }, + "nperseg": 16384, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d" + }, + "pa_evaluation_device": "cuda:0", + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 1000, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 640000000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 2.295543024968947, + "task": "dpd_ila", + "training_relative_residual": 0.01053396426141262 + }, + "gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_GRU_H_16_F_200_P_994.csv", + "sha256": "0d5e9c035a56606e6eeaedd6752fc3e760b3c3490f6739d57c471375a84b3cde", + "size_bytes": 424, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_GRU_H_16_F_200_P_994.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_GRU_H_16_F_200_P_994.pt", + "sha256": "a7764785b413476b59d5f5e9528de2c8784b9abfa11f36a89383d836452dda9b", + "size_bytes": 6832, + "source_path": "save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_GRU_H_16_F_200_P_994.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_GRU_H_16_F_200_P_994.csv", + "sha256": "0d5ab7a838ed9b7b2442d04efa0770da9c7ff7c707b67739d6ada72cc1593902", + "size_bytes": 62460, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_GRU_H_16_F_200_P_994.csv" + } + }, + "completed_epochs": 300, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -57.26420765, + "aclr_left_db": -57.675965, + "aclr_right_db": -56.85245031, + "evm_db": -53.97904848, + "loss": 2.95e-06, + "nmse_db": -49.760174 + }, + "validation": { + "aclr_avg_db": -57.35296233, + "aclr_left_db": -57.60844419, + "aclr_right_db": -57.09748047, + "evm_db": -53.73140247, + "loss": 2.9e-06, + "nmse_db": -49.905857 + } + }, + "model": { + "backbone": "gru", + "display_name": "GRU-H16", + "hidden_size": 16, + "model_id": "DPD_S_0_M_GRU_H_16_F_200_P_994", + "parameters": 994 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "DPA_160MHz", + "dpd_model_key": "gru", + "pa_model_key": "tres_gru", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d" + }, + "selected_epoch_zero_based": 297, + "selected_learning_rate": 5e-05, + "task": "dpd_training", + "train_loss": 1.097e-05 + }, + "mp": { + "artifact": { + "path": "benchmark_report_dpa_dpd_mp.json", + "sha256": "3a1353f590f28c9c4a058315fcedade248f55976912762e142e3fed6d5a6da57", + "size_bytes": 2908 + }, + "basis_configuration": { + "K": 5, + "Q": 100 + }, + "coefficient_l2_norm": 3.943410873413086, + "column_scale_max": 210.84939575195312, + "column_scale_min": 38.87131118774414, + "column_scaling": "l2", + "complex_coefficients": 500, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "indirect_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -50.9351411277973, + "aclr_left_db": -49.44417246777463, + "aclr_right_db": -52.426109787819975, + "evm_db": -50.50021219066906, + "nmse_db": -43.193756103515625 + }, + "validation": { + "aclr_avg_db": -50.63197307642169, + "aclr_left_db": -49.06701086616408, + "aclr_right_db": -52.19693528667929, + "evm_db": -50.42040796343259, + "nmse_db": -43.55812072753906 + } + }, + "model": { + "backbone": "mp", + "display_name": "MP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 1000 + }, + "nperseg": 16384, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d" + }, + "pa_evaluation_device": "cuda:0", + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 1000, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 640000000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 2.295543024968947, + "task": "dpd_ila", + "training_relative_residual": 0.010844217613339424 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv", + "sha256": "58a96914a057b9e4a58468d8a84f65c435e6f65da932c94fe03a2ab0b2a1d1ed", + "size_bytes": 428, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.pt", + "sha256": "751b9a9228713a3a25fc5bef2ca6e45e6a4e5bbd979651ea5666657b03fa5062", + "size_bytes": 7541, + "source_path": "save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv", + "sha256": "d685adb6158f233b9f77f34a5d2375f66caadccbc65a49e1c6ebede68abde8d2", + "size_bytes": 63950, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_GRU_H_15_F_200_P_999.csv" + } + }, + "completed_epochs": 300, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -60.37169906, + "aclr_left_db": -60.27122719, + "aclr_right_db": -60.47217092, + "evm_db": -57.84510577, + "loss": 1.32e-06, + "nmse_db": -53.19377 + }, + "validation": { + "aclr_avg_db": -60.74250569, + "aclr_left_db": -60.45794513, + "aclr_right_db": -61.02706624, + "evm_db": -57.14251367, + "loss": 1.44e-06, + "nmse_db": -52.910625 + } + }, + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H15", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_GRU_H_15_F_200_P_999", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "DPA_160MHz", + "dpd_model_key": "tres_gru", + "pa_model_key": "tres_gru", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d" + }, + "selected_epoch_zero_based": 296, + "selected_learning_rate": 5e-05, + "task": "dpd_training", + "train_loss": 1.49e-06 + } + }, + "dpd_tres_deltagru_by_pa_surrogate": { + "tres_deltagru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "8c583529041e534dc5cb1fbcd54c06d12f7578342c12309684da7c70521ef7d7", + "size_bytes": 464, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt", + "sha256": "a3cdb6bdf7c2fccb9a0bb7daa8fad29c659632bc6355287ad60fe199e3197e2e", + "size_bytes": 8160, + "source_path": "save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "dd135e3d4951358d37f84799b2b4752f92da9e4fb5798a34b7ff4ddddcfa5d2d", + "size_bytes": 72065, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_DELTAGRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -61.65856766, + "aclr_left_db": -61.4514239, + "aclr_right_db": -61.86571143, + "evm_db": -59.39778461, + "loss": 1e-06, + "nmse_db": -54.372208 + }, + "validation": { + "aclr_avg_db": -61.92533611, + "aclr_left_db": -61.54620906, + "aclr_right_db": -62.30446316, + "evm_db": -58.24780618, + "loss": 1.11e-06, + "nmse_db": -54.115932 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "dataset": "DPA_160MHz", + "dpd_model_key": "tres_deltagru", + "pa_model_key": "tres_deltagru", + "sha256": "980f4f8eb48da7d96bf5b0ce6b348755398263cc34a83fea8209dec45d0d45a7" + }, + "pa_surrogate": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "980f4f8eb48da7d96bf5b0ce6b348755398263cc34a83fea8209dec45d0d45a7", + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H27 (THX=THH=0)", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "model_key": "tres_deltagru", + "selection_metric": "validation NMSE" + }, + "selected_epoch_zero_based": 293, + "selected_learning_rate": 0.00015625, + "task": "dpd_training", + "train_loss": 1.15e-06 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "549232530b4ddba266549b6ec4d6d36b329998bf625b8763fd20f91ee0cd6d4f", + "size_bytes": 464, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/best/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt", + "sha256": "f69815ba379e01c1f492e906af021282d6045854ae1860e6c1b6aae0a611b64f", + "size_bytes": 8160, + "source_path": "save/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv", + "sha256": "f0614b1f2a87acd1a0573388af81e71131aacea773404826bd80f24c29052f04", + "size_bytes": 72105, + "source_path": "log/DPA_160MHz/train_dpd/PA_S_0_M_TRES_GRU_H_27_F_200/history/DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "direct_learning_architecture", + "metrics": { + "test": { + "aclr_avg_db": -61.04416081, + "aclr_left_db": -60.57591376, + "aclr_right_db": -61.51240785, + "evm_db": -58.84856644, + "loss": 1.03e-06, + "nmse_db": -54.244564 + }, + "validation": { + "aclr_avg_db": -61.54248186, + "aclr_left_db": -61.01506389, + "aclr_right_db": -62.06989982, + "evm_db": -57.93123352, + "loss": 1.15e-06, + "nmse_db": -53.868393 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000", + "parameters": 999 + }, + "pa_checkpoint_binding": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "dataset": "DPA_160MHz", + "dpd_model_key": "tres_deltagru", + "pa_model_key": "tres_gru", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d" + }, + "pa_surrogate": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d", + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H27", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "model_key": "tres_gru", + "selection_metric": "validation NMSE" + }, + "selected_epoch_zero_based": 295, + "selected_learning_rate": 7.813e-05, + "task": "dpd_training", + "train_loss": 1.29e-06 + } + }, + "pa_models": { + "gmp": { + "artifact": { + "path": "benchmark_report_dpa_pa_gmp.json", + "sha256": "070485cf88eebe45fbb0337976e315a03c31b56b1c50ace22522d4803506f885", + "size_bytes": 2818 + }, + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 30, + "Lb": 30, + "Lc": 30, + "Mb": 5, + "Mc": 5 + }, + "coefficient_l2_norm": 10.034621238708496, + "column_scale_max": 165.83279418945312, + "column_scale_min": 7.7007365226745605, + "column_scaling": "l2", + "complex_coefficients": 1350, + "condition_number": 1985705.7386572887, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": false, + "least_squares_rank": 918, + "method": "direct_least_squares", + "metrics": { + "test": { + "aclr_avg_db": -34.80149418436232, + "aclr_left_db": -35.08585122730327, + "aclr_right_db": -34.517137141421365, + "evm_db": -42.38386821531598, + "nmse_db": -39.40143419744728 + }, + "validation": { + "aclr_avg_db": -34.63113398986266, + "aclr_left_db": -34.9108251117308, + "aclr_right_db": -34.35144286799453, + "evm_db": -41.65365023479418, + "nmse_db": -38.908763378736715 + } + }, + "model": { + "backbone": "gmp", + "display_name": "GMP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 2700 + }, + "nperseg": 16384, + "pa_checkpoint_binding": null, + "pa_evaluation_device": null, + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 2700, + "regularization": "truncated_svd", + "retained_singular_value_min": 0.0011493448400869966, + "sample_rate_hz": 640000000.0, + "singular_value_cutoff": 0.0011460232734680177, + "singular_value_max": 11.460232734680176, + "singular_value_min": 5.771365067630541e-06, + "solver": "torch.linalg.svd", + "solver_driver": "gesvdj", + "solver_mode": "truncated_svd", + "svd_rcond": 0.0001, + "target_gain": 2.295543024968947, + "task": "pa_modeling", + "training_relative_residual": 0.010578107088804245 + }, + "gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/best/PA_S_0_M_GRU_H_28_F_200_P_2746.csv", + "sha256": "7c5be39701ff59e16bdb4e23a4ad714adfdb611b6dd5a09b9079ae600016f6b3", + "size_bytes": 421, + "source_path": "log/DPA_160MHz/train_pa/best/PA_S_0_M_GRU_H_28_F_200_P_2746.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_pa/PA_S_0_M_GRU_H_28_F_200_P_2746.pt", + "sha256": "c1d1dc696df19d7519def6dda23a22fcc82589224d43e5caef85cd84a573851f", + "size_bytes": 13808, + "source_path": "save/DPA_160MHz/train_pa/PA_S_0_M_GRU_H_28_F_200_P_2746.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/history/PA_S_0_M_GRU_H_28_F_200_P_2746.csv", + "sha256": "5c8a49387a5b08af6eb09f044315011142dd642fd780277bac36f031ccedd95e", + "size_bytes": 62434, + "source_path": "log/DPA_160MHz/train_pa/history/PA_S_0_M_GRU_H_28_F_200_P_2746.csv" + } + }, + "completed_epochs": 300, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -34.83799881, + "aclr_left_db": -35.03673896, + "aclr_right_db": -34.63925866, + "evm_db": -42.82822831, + "loss": 4.992e-05, + "nmse_db": -39.49081 + }, + "validation": { + "aclr_avg_db": -34.76237982, + "aclr_left_db": -34.95988762, + "aclr_right_db": -34.56487201, + "evm_db": -42.19929651, + "loss": 5.342e-05, + "nmse_db": -39.080143 + } + }, + "model": { + "backbone": "gru", + "display_name": "GRU-H28", + "hidden_size": 28, + "model_id": "PA_S_0_M_GRU_H_28_F_200_P_2746", + "parameters": 2746 + }, + "selected_epoch_zero_based": 34, + "selected_learning_rate": 0.00125, + "task": "pa_modeling", + "train_loss": 7.127e-05 + }, + "mp": { + "artifact": { + "path": "benchmark_report_dpa_pa_mp.json", + "sha256": "a5a6cb62b210f6e22f4107ec3f17338aafc3ac351c54b8ebda6dd88620c1ba0f", + "size_bytes": 2539 + }, + "basis_configuration": { + "K": 9, + "Q": 150 + }, + "coefficient_l2_norm": 1562.5501708984375, + "column_scale_max": 165.83279418945312, + "column_scale_min": 8.192251205444336, + "column_scaling": "l2", + "complex_coefficients": 1350, + "condition_number": null, + "device": "cuda:0", + "dtype": "torch.complex64", + "full_rank_assumption": true, + "least_squares_rank": null, + "method": "direct_least_squares", + "metrics": { + "test": { + "aclr_avg_db": -34.89261561507139, + "aclr_left_db": -35.144338147000205, + "aclr_right_db": -34.640893083142586, + "evm_db": -40.903476587562324, + "nmse_db": -38.350608084028664 + }, + "validation": { + "aclr_avg_db": -34.70183472040674, + "aclr_left_db": -34.931857098557934, + "aclr_right_db": -34.47181234225553, + "evm_db": -40.39402625607741, + "nmse_db": -37.98229716673007 + } + }, + "model": { + "backbone": "mp", + "display_name": "MP", + "parameter_convention": "two real degrees of freedom per complex coefficient", + "parameters": 2700 + }, + "nperseg": 16384, + "pa_checkpoint_binding": null, + "pa_evaluation_device": null, + "parameter_count_convention": "two real degrees of freedom per complex coefficient", + "real_parameters": 2700, + "regularization": null, + "retained_singular_value_min": null, + "sample_rate_hz": 640000000.0, + "singular_value_cutoff": null, + "singular_value_max": null, + "singular_value_min": null, + "solver": "torch.linalg.lstsq", + "solver_driver": "gels", + "solver_mode": "gels", + "svd_rcond": null, + "target_gain": 2.295543024968947, + "task": "pa_modeling", + "training_relative_residual": 0.011598878540098667 + }, + "tres_deltagru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/best/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv", + "sha256": "cce446f3d8d7d7d49b9c9c72f5cb9091d53c5c7c70637ed42335d3bdfbc45b90", + "size_bytes": 462, + "source_path": "log/DPA_160MHz/train_pa/best/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt", + "sha256": "980f4f8eb48da7d96bf5b0ce6b348755398263cc34a83fea8209dec45d0d45a7", + "size_bytes": 14596, + "source_path": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/history/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv", + "sha256": "d315bfcd3c3f823a44a5b0e7974191139c3315cf20535b90f6696dd4b55bbadf", + "size_bytes": 72059, + "source_path": "log/DPA_160MHz/train_pa/history/PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751.csv" + } + }, + "completed_epochs": 300, + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -34.89409696, + "aclr_left_db": -35.14318315, + "aclr_right_db": -34.64501077, + "evm_db": -43.31956339, + "loss": 4.677e-05, + "nmse_db": -39.735275 + }, + "validation": { + "aclr_avg_db": -34.70573261, + "aclr_left_db": -34.96491134, + "aclr_right_db": -34.44655387, + "evm_db": -42.96421325, + "loss": 4.832e-05, + "nmse_db": -39.480164 + } + }, + "model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H27 (THX=THH=0)", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_DELTAGRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "selected_epoch_zero_based": 16, + "selected_learning_rate": 0.0025, + "task": "pa_modeling", + "train_loss": 5.689e-05 + }, + "tres_gru": { + "artifacts": { + "best_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/best/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv", + "sha256": "b6abe61b2433d6d170b1c2483b97d3275e6f93642053baf894de4a25a69afc8c", + "size_bytes": 425, + "source_path": "log/DPA_160MHz/train_pa/best/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv" + }, + "checkpoint": { + "archived_path": "artifacts/save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d", + "size_bytes": 14773, + "source_path": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt" + }, + "history_csv": { + "archived_path": "artifacts/log/DPA_160MHz/train_pa/history/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv", + "sha256": "19102ffefe250224113d0973e5a0e9cd1bb2de8375e9f93641c85c13faea3861", + "size_bytes": 63971, + "source_path": "log/DPA_160MHz/train_pa/history/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.csv" + } + }, + "completed_epochs": 300, + "method": "supervised_training", + "metrics": { + "test": { + "aclr_avg_db": -34.90779123, + "aclr_left_db": -35.15503937, + "aclr_right_db": -34.66054308, + "evm_db": -43.1446587, + "loss": 4.741e-05, + "nmse_db": -39.70443 + }, + "validation": { + "aclr_avg_db": -34.72994139, + "aclr_left_db": -34.99051361, + "aclr_right_db": -34.46936918, + "evm_db": -42.58117816, + "loss": 5.033e-05, + "nmse_db": -39.32285 + } + }, + "model": { + "backbone": "tres_gru", + "display_name": "TRes-GRU-H27", + "hidden_size": 27, + "model_id": "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751", + "parameters": 2751 + }, + "selected_epoch_zero_based": 25, + "selected_learning_rate": 0.00125, + "task": "pa_modeling", + "train_loss": 5.561e-05 + } + }, + "reference_pa": { + "checkpoint": "save/DPA_160MHz/train_pa/PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt", + "checkpoint_sha256": "e09da8eebdbf35746975ebf5013f6eb82d02adea3c9545bd50fe88a2bfe6ce6d", + "model_key": "tres_gru", + "selection_metric": "validation NMSE" + }, + "signal": { + "dataset_file_sha256": { + "datasets/DPA_160MHz/test_input.csv": "7fb259816af61b34ce6a966858787ca164fd1aed287fafa6005c91ae4abcd5fb", + "datasets/DPA_160MHz/test_output.csv": "06eab1afc94c0ee3140cabc785ed1a81e7e3446d0a580aebc749c5853a1d07c5", + "datasets/DPA_160MHz/train_input.csv": "3de4709bcb4e9e5c155c075ab8e210de3b8545f73a43edaf05885b41a849a9d1", + "datasets/DPA_160MHz/train_output.csv": "e2d852dccac4f0db8934746f8c14aedb316edd5589c972d7a10d6dc85c0775c2", + "datasets/DPA_160MHz/val_input.csv": "ad04f3f272d6974429c7a5b6ee6bbcca91c180157066a5994ae7bb5912053c97", + "datasets/DPA_160MHz/val_output.csv": "efdfeb3d261a284024daa53378201917d1d7af96e128fa65a5c44c5266665c15" + }, + "papr": { + "absolute_peak_db": 10.139761791608466, + "ccdf_db": 10.131554833438559, + "ccdf_probability": 1e-05, + "quantile_method": "numpy linear" + }, + "samples": 491520, + "target_gain": 2.2955430249689464 + }, + "spec_sha256": "4e777ff275bd1b2676c5ffc6a34d3432b56aedb3ed78935734485f3f9c2f1e8b" + } + }, + "derived_artifacts": { + "benchmark_delta_dpd_results.png": { + "sha256": "cb2f7696c4654b0b735137c4498044edf756f10c6f57db6badb693899e8cdfa6", + "size_bytes": 112421 + }, + "benchmark_report.md": { + "sha256": "edd5a017414006e7b678e293ae1b7e3d61390b898e28e9e9f199e34a7c5a98d7", + "size_bytes": 13602 + }, + "benchmark_results.png": { + "sha256": "6bf8c8f61ba5664e6fd26ee849c6ac1718b8e238e1aff92fee7fe028379219e7", + "size_bytes": 201177 + } + }, + "environment": { + "cuda": "13.2", + "cudnn": 92000, + "gpu": "NVIDIA RTX PRO 6000 Blackwell Workstation Edition", + "gpu_device_index": 0, + "numpy": "2.5.1", + "pandas": "3.0.5", + "platform": "Linux-7.0.0-15-generic-x86_64-with-glibc2.43", + "python": "3.13.14", + "scipy": "1.18.0", + "torch": "2.13.0+cu132", + "training_cudnn_benchmark": true, + "training_deterministic_algorithms": false, + "training_reproducibility_level": "soft" + }, + "generated_at_utc": "2026-07-26T20:45:35.368238+00:00", + "git": { + "branch": "benchmark-fix", + "commit": "3df35e081e6e41463fa46f21778c72a748823274", + "dirty": true, + "status": [ + "M .gitignore", + " M arguments.py", + " D benchmark/BENCHMARK_FIX_PLAN.md", + " D benchmark/benchmark_end_to_end.py", + " M benchmark/benchmark_report.md", + " D benchmark/benchmark_training_speed.py", + " D benchmark/benchmark_training_speed_report.md", + " D benchmark/benchmark_tres_deltagru.py", + " D benchmark/benchmark_volterra_qr.py", + " D benchmark/pa_modeling_investigation.md", + " D benchmark/report/artifact.json", + " D benchmark/results/e2e_dpa200_stride16_comparison_eager.json", + " D benchmark/results/e2e_dpa200_stride16_comparison_fused.json", + " D benchmark/results/e2e_dpa200_stride16_comparison_graph.json", + " D benchmark/results/e2e_dpa200_stride16_eager_optimized/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride16_eager_optimized/outputs.npz", + " D benchmark/results/e2e_dpa200_stride16_eager_optimized/result.json", + " D benchmark/results/e2e_dpa200_stride16_fused_optimized/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride16_fused_optimized/outputs.npz", + " D benchmark/results/e2e_dpa200_stride16_fused_optimized/result.json", + " D benchmark/results/e2e_dpa200_stride16_graph_optimized/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride16_graph_optimized/outputs.npz", + " D benchmark/results/e2e_dpa200_stride16_graph_optimized/result.json", + " D benchmark/results/e2e_dpa200_stride16_original/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride16_original/outputs.npz", + " D benchmark/results/e2e_dpa200_stride16_original/result.json", + " D benchmark/results/e2e_dpa200_stride1_comparison_graph.json", + " D benchmark/results/e2e_dpa200_stride1_comparison_graph_repeat.json", + " D benchmark/results/e2e_dpa200_stride1_graph_optimized/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride1_graph_optimized/outputs.npz", + " D benchmark/results/e2e_dpa200_stride1_graph_optimized/result.json", + " D benchmark/results/e2e_dpa200_stride1_graph_repeat/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride1_graph_repeat/outputs.npz", + " D benchmark/results/e2e_dpa200_stride1_graph_repeat/result.json", + " D benchmark/results/e2e_dpa200_stride1_original/batch_orders.npz", + " D benchmark/results/e2e_dpa200_stride1_original/outputs.npz", + " D benchmark/results/e2e_dpa200_stride1_original/result.json", + " D benchmark/results/end_to_end_summary.json", + " D benchmark/results/tres_deltagru_rtx4090_laptop.json", + " D benchmark/train_all_pa_2700.sh", + " D benchmark/tres_deltagru_end_to_end_report.html", + " D benchmark/tres_deltagru_training_speed.md", + " M modules/paths.py", + " M modules/train_funcs.py", + " M opendpd/api.py", + " M project.py", + " M steps/run_dpd.py", + " M steps/train_dpd.py", + " M steps/train_pa.py", + " M tests/test_api.py", + " M tests/test_cli.py", + " M tests/test_tres_deltagru_optimized.py", + "?? benchmark/benchmark_delta_dpd_results.png", + "?? benchmark/benchmark_results.png", + "?? benchmark/benchmark_volterra.py", + "?? benchmark/collect_benchmark_report.py", + "?? benchmark/reproduce_benchmark_report.sh", + "?? benchmark/results/benchmark_report_results.json", + "?? tests/test_benchmark_regressions.py", + "?? tests/test_metric_configuration.py" + ] + }, + "metric_note": "EVM is the repository-specific normalized mean absolute complex-spectrum error across configured main-channel subchannels, not demodulated constellation EVM. DPD metrics are simulated through a frozen learned PA surrogate.", + "publication": { + "all_jobs_exit_zero": true, + "note": "The canonical manifest contains the published model subset. Source-run hashes, job totals, and archives cover the complete source run and can include experiments outside that subset; source-archive members were verified against the pre-run source hash map.", + "original_collector_manifest_sha256": "65daa8b3f91d51f3ecb200e0477e44d20e0e23bd8201afbdb78087aa20a2e73b", + "original_schema_version": 6, + "published_dpd_pa_surrogates": [ + "tres_gru", + "tres_deltagru" + ], + "published_models": [ + "mp", + "gmp", + "gru", + "tres_gru", + "tres_deltagru" + ], + "published_models_by_stage": { + "dpd": [ + "mp", + "gmp", + "gru", + "tres_gru" + ], + "dpd_pa_surrogate_sensitivity": [ + "tres_deltagru" + ], + "pa_modeling": [ + "mp", + "gmp", + "gru", + "tres_gru", + "tres_deltagru" + ] + }, + "published_result_count": 22, + "source_job_count": 24, + "source_result_count": 22, + "source_run_id": "20260726T161245Z", + "source_snapshot_created_during_run_utc": "2026-07-26T16:12:47.861423+00:00", + "source_snapshot_member_hashes_match_pre_run_context": true, + "total_job_duration_seconds": 16369 + }, + "recipes": { + "neural_dpd": { + "batch_size": 64, + "decay_factor": 0.5, + "epochs": 300, + "frame_length": 200, + "frame_stride": 1, + "initial_learning_rate": 0.005, + "loss": "MSE", + "lr_schedule": "ReduceLROnPlateau", + "minimum_learning_rate": 5e-05, + "optimizer": "AdamW", + "optimizer_betas": [ + 0.9, + 0.999 + ], + "optimizer_eps": 1e-08, + "optimizer_weight_decay": 0.01, + "patience": 5, + "scheduler_cooldown": 0, + "scheduler_eps": 1e-08, + "scheduler_threshold": 0.0001, + "scheduler_threshold_mode": "rel", + "seed": 0 + }, + "neural_dpd_pa_surrogate_sensitivity": { + "batch_size": 64, + "decay_factor": 0.5, + "dpd_model": { + "backbone": "tres_deltagru", + "delta_thresholds": { + "hidden": 0.0, + "input": 0.0 + }, + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "hidden_size": 15, + "model_id": "DPD_S_0_M_TRES_DELTAGRU_H_15_F_200_P_999_THX_0.000_THH_0.000", + "parameters": 999 + }, + "epochs": 300, + "frame_length": 200, + "frame_stride": 1, + "initial_learning_rate": 0.005, + "loss": "MSE", + "lr_schedule": "ReduceLROnPlateau", + "minimum_learning_rate": 5e-05, + "optimizer": "AdamW", + "optimizer_betas": [ + 0.9, + 0.999 + ], + "optimizer_eps": 1e-08, + "optimizer_weight_decay": 0.01, + "pa_surrogates": [ + "tres_gru", + "tres_deltagru" + ], + "patience": 5, + "scheduler_cooldown": 0, + "scheduler_eps": 1e-08, + "scheduler_threshold": 0.0001, + "scheduler_threshold_mode": "rel", + "seed": 0 + }, + "neural_pa": { + "batch_size": 64, + "decay_factor": 0.5, + "epochs": 300, + "frame_length": 200, + "frame_stride": 1, + "initial_learning_rate": 0.005, + "loss": "MSE", + "lr_schedule": "ReduceLROnPlateau", + "minimum_learning_rate": 5e-05, + "optimizer": "AdamW", + "optimizer_betas": [ + 0.9, + 0.999 + ], + "optimizer_eps": 1e-08, + "optimizer_weight_decay": 0.01, + "patience": 5, + "scheduler_cooldown": 0, + "scheduler_eps": 1e-08, + "scheduler_threshold": 0.0001, + "scheduler_threshold_mode": "rel", + "seed": 0 + }, + "polynomial_dpd": { + "column_scaling": "l2", + "method": "indirect_learning_architecture", + "models": { + "gmp": { + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 20, + "Lb": 20, + "Lc": 20, + "Mb": 3, + "Mc": 2 + }, + "complex_coefficients": 500, + "real_parameters": 1000 + }, + "mp": { + "basis_configuration": { + "K": 5, + "Q": 100 + }, + "complex_coefficients": 500, + "real_parameters": 1000 + } + }, + "parameter_convention": "two real degrees of freedom per complex coefficient", + "solver": "torch.linalg.lstsq", + "solver_driver": "gels" + }, + "polynomial_pa": { + "column_scaling": "l2", + "method": "direct_least_squares", + "models": { + "gmp": { + "basis_configuration": { + "Ka": 5, + "Kb": 4, + "Kc": 4, + "La": 30, + "Lb": 30, + "Lc": 30, + "Mb": 5, + "Mc": 5 + }, + "complex_coefficients": 1350, + "real_parameters": 2700 + }, + "mp": { + "basis_configuration": { + "K": 9, + "Q": 150 + }, + "complex_coefficients": 1350, + "real_parameters": 2700 + } + }, + "parameter_convention": "two real degrees of freedom per complex coefficient", + "solvers": { + "gmp": { + "driver": "gesvdj", + "implementation": "torch.linalg.svd", + "mode": "truncated_svd", + "regularization": { + "relative_cutoff": 0.0001, + "type": "truncated_svd" + } + }, + "mp": { + "driver": "gels", + "implementation": "torch.linalg.lstsq", + "mode": "gels", + "regularization": null + } + } + } + }, + "run_file_sha256": { + "commands.log": "cc18c355294a0e2d9bb7056e8c06bbfd12492664eb3968976bfe913090ceb1b0", + "context_before.json": "9d578fbc05fa1d943193d5397dd6741c65b1658baaaea9f773a2d50ccfe20808", + "git_branch_before.txt": "758faa270e2de48355d38d5ec23cc5b5d89f668c8cd8b68c0a3b10af638fbb97", + "git_commit_before.txt": "c3001cec76c14bb641eeb21bbbf88dd42bd33fe6d73be1dc3bb1759dc646e98c", + "git_status_before.txt": "e6225bec19c0425e36bb3bb9bd49e2347500bf842b9f229f02d090538a4ece7b", + "jobs.tsv": "d580289f06f1be5f571bc9a8467910a41a0c0eb839834c6683b7468947cfa66c", + "pa_checkpoint_bindings.tsv": "cbbae37c50cfeb831d338e75b52bec9653657d90f4ae5544475c87e4f67f3770", + "recipe.json": "aa372944dbd085fca23f859edd1de5df3800282867f5263696c14d3e57dc991c", + "source_snapshot.tar.gz": "7ecc3ae72750b24ece1e2aa0814f539a7cf672b11da2f37d31e67f7a517f4811" + }, + "schema_version": 6, + "source_file_sha256": { + "arguments.py": "2d00505312ab0d87d58dcebaebaf3ab15d77b1c29e79ced7c8ec2bda0f9e7f6d", + "backbones/__init__.py": "c36ccba9afef9c5c2aad3776939a9746379c435054afd8673fc6096ad3e6d12c", + "backbones/apnrru.py": "9ec899243bc04a497cf479c49dea8bd875e50884c7cc0157581828fdde2a4e59", + "backbones/bojanet.py": "f1ae8fc298935afe9764836bda769cf31437cc325ddc1141927153dab68fc3c2", + "backbones/cuda_graph_frozen_dgru.py": "033e326dfa4782cd13e0c37f62d0581d8a8513a536b39130bfc9d3ee69b2682c", + "backbones/deltagru.py": "7582250e17a4bb5b27d3244bf06145388a0506c3e83205db6df79d259f6ed4e1", + "backbones/deltajanet.py": "7f2a48ffb375d403434aace161f858319ae981e90b7828ec5e665c0d6d9ec950", + "backbones/dgru.py": "6bcb7dcff9abba5424c1312515570eef3614b4c62bb7f8c228772a467d2720ef", + "backbones/dvrjanet.py": "e7112cbc194b1b95bc96b84075457a17ef7671cab38279b8086c65297873aee4", + "backbones/gmp.py": "53339dc48a8da7333a11756996223a4921759298604159ad7f2499e0c945585c", + "backbones/gru.py": "9e1e17ac9ad936bf91c821703a9137260e9df352541155262a5c0a8c0e1a6e30", + "backbones/lstm.py": "7cc761002c83f09f2e38878547a3c065c01fbfb64ee4db9c1cbe8510d17c1bed", + "backbones/mcldnn.py": "1cc9af20c6c35f4efe5fec95375b077a4e10563a4b5c2fa501cc3dde8940c464", + "backbones/pgjanet.py": "672f20e50e920826e570640efe6172e9f02d5e35eb5b66a607beddacf9ebfbfb", + "backbones/qgru.py": "5b029c99e8145628834dd0cc88cd68b9d1c9e1ef3174c8a85a8d46b68ccf6536", + "backbones/qgru_amp1.py": "21296ba3777722f5ff7d0ee9c0a1dfcd8d31c1da65c836d9577bdd10d49d245c", + "backbones/rvtdcnn.py": "d98b2a3e2161e4e74ebc37dc65f5cdd0a247644e363393d57fc4bc245e001923", + "backbones/tcn.py": "3d3c7048b36591409d90a7b0e4c4fb7ea73a01018ef73192e42df252c66cd827", + "backbones/tres_deltagru.py": "0351d322ce3ae615d0f1adfde2b76e8afaa0d85d6f9dc92f511905eca7762925", + "backbones/tres_gru.py": "0c03c302a7e98751f654f0e56da86498b9fc01723717c01e83dd98571be1d1e6", + "backbones/triton_deltagru.py": "1404d1bb55b512164eac1f281f318d7744cb682ee95fc25c170d6d5d1a9f4718", + "backbones/vdlstm.py": "bf268d65f47ecd18c31e68aaadcb3a26a433d7483b0fe888e39ef89f33d05619", + "benchmark/benchmark_report.md": "8d1b0ec5b368da42d3b35f4b71f3c31aafe14542e4167de793c8b73856c062c3", + "benchmark/benchmark_volterra.py": "79cb35031113908fdd177174c34383dad381f8e48b4eab008b4c14ef5988285b", + "benchmark/collect_benchmark_report.py": "f75caaa35812e8cf15635af523627e14b5740bddfef3c02535262c25aa81be95", + "benchmark/reproduce_benchmark_report.sh": "266949c88a18f4b7072c0809d0900bcd7abe3ac9584215b41fba932d3d162654", + "main.py": "07194242bd0b7fa8d1d05f1661daee42fc668ccea72fed04e661942e15239c26", + "models.py": "e1a7887a5e8621620088046e8aac4429247f039cddfedcb603b01fee57a979f7", + "modules/__init__.py": "c36ccba9afef9c5c2aad3776939a9746379c435054afd8673fc6096ad3e6d12c", + "modules/cuda_graph_training.py": "dc6d9c436eb8a06145f9e5472cd6cb72deebcc2b08f44a278a3be21409dde620", + "modules/data_collector.py": "7f3b55f38f80820c10d399d161b8cd233bfaf8837dc43613dd1410d6ed649a4a", + "modules/loggers.py": "192d4851f41581e211a32430856e630cb3e92fcb88e73525ec555acb0a76b3d6", + "modules/paths.py": "43c3cec2727db9c0b6206831440e9a26992a1366d151ab025ef2f35536403d66", + "modules/train_funcs.py": "c3c6057f8822ccd470b626e775fb66e6116954f4185265523f05146fb5ab2432", + "opendpd/__init__.py": "5535b674cf7a2b92563577dc5dc0036301e8a8ec06e10e718f858b976c7f1fee", + "opendpd/api.py": "e68edf808ee2feaa17c9563bc8adbb5b956e097649f3382449af31a308f8a66c", + "opendpd/cli.py": "1142070bd94de6dfe61e241f66391ee4f2779bae215d030d2143aa59f07997d1", + "project.py": "759a7c4d1d1d73a4444e83b83269ce346a3b1a849048c8ed8fc21ddc751098df", + "quant/__init__.py": "77870e58c91b6f96aa99dcfbd51cbc2af0647bc5471f6f1bce75fe3dc9844366", + "quant/modules/__init__.py": "5041f5f1c0c70267661470188147ee9c59b88dc75b1c9ba43a0f64f6c63a5929", + "quant/modules/gru.py": "768312fd67030bb909f3340ea61da470a8c6092094e2a161cd6fbb544f9416d4", + "quant/modules/ops.py": "2766edd66d3b4578c6c9b8c741950914de1f189db497351caf49ba5d00846eb9", + "quant/qmodules/__init__.py": "c65bf92be624ac36c266227fd10b217652ab4e188c4ac12b1c3d4c1a318ef1ad", + "quant/qmodules/quant_layers.py": "24352078f035768a8d640550f68b5f44048fa1f7d2cbaae09f50376fdd36a54c", + "quant/qmodules/quant_ops.py": "99ec701d2b7d095f3e75114e59a0843fdda849cf2ea82abdee096e67c822e6ac", + "quant/qmodules/quantizers.py": "1c44adfef579bf97afb314dd10074ba6eacc5207349873665b4e9e1ba2719e8b", + "quant/quant_envs.py": "803436d91459e76c0530951f1f068fe85385abac753a233b4d43a4553ade70c5", + "quant/utlis/__init__.py": "ce2b4d895b803d84d02351b68a1ebf536ebda3de11bb621637aa478610dd5973", + "quant/utlis/hooks.py": "4cf0063c4aa4c007b5f448d7427a7491e09793ed985a0337c1f170f452e0b3fc", + "quant/utlis/measures.py": "6b06453d2e8e68db39262eff3fe724764dcf5f8d15af1f4f8ebfdfa844b19649", + "setup.py": "ee47c04e451de6cf7fb05290cf8535a8311dc202ebaeedd41ae767b2513668cc", + "steps/__init__.py": "c36ccba9afef9c5c2aad3776939a9746379c435054afd8673fc6096ad3e6d12c", + "steps/plot.py": "0bd2c05b45df599d4d36f6393ada2454344f4bce84bbdadb6c586faee0a0bcdb", + "steps/run_dpd.py": "3d6dacc3fc11f17de6622a7ceadca6033a0fad279a86113eff61a7e8c08891be", + "steps/train_dpd.py": "03326bb221271531eb2db355604d0174d554f67d91211a9e32f5811760775f4f", + "steps/train_pa.py": "689fe78e150fa6600dca31f60d7b4f7493eb3f9b879fb4beabeb25b9ac381c1c", + "test_installation.py": "2fa9123da50e24dadc8bbe659440e7a2f9bd11d2482523ac1a323ffcfb490b0f", + "utils/__init__.py": "c36ccba9afef9c5c2aad3776939a9746379c435054afd8673fc6096ad3e6d12c", + "utils/metrics.py": "19e525834c96371e9837f909b8e848d4b32ff49ef52855e3281ec3f7682cd992", + "utils/plotting.py": "2ac10c110559b7ed999647ecdd9f06bde50044a8a1c807d33876f91729deedcf", + "utils/split_dataset.py": "472a9956db640dad7fb24e605d7722796d067366081d4e48d06a6e96f65ccc03", + "utils/util.py": "dcfe263cf2b3bb55e18f13439c3aed196a49d975ed3bfbab81f781a4cbcfef8b" + }, + "source_snapshot": { + "file_count": 62, + "path": "source_snapshot.tar.gz", + "sha256": "7ecc3ae72750b24ece1e2aa0814f539a7cf672b11da2f37d31e67f7a517f4811" + }, + "started_at_utc": "2026-07-26T16:12:47.178818+00:00" +} diff --git a/models.py b/models.py index 76da3b7..a53bcd4 100644 --- a/models.py +++ b/models.py @@ -126,6 +126,13 @@ def __init__(self, input_size, hidden_size, num_layers, backbone_type, window_si thx=self.thx, thh=self.thh, bias=self.bias) + elif backbone_type == 'tres_gru': + from backbones.tres_gru import TResGRU + self.backbone = TResGRU(input_size=6, + hidden_size=self.hidden_size, + output_size=self.output_size, + num_layers=self.num_layers, + bias=self.bias) elif backbone_type == 'tcn': from backbones.tcn import TCN self.backbone = TCN(hidden_channels=self.hidden_size) @@ -147,11 +154,14 @@ def __init__(self, input_size, hidden_size, num_layers, backbone_type, window_si pass def forward(self, x, h_0=None): - device = x.device batch_size = x.size(0) # NOTE: dim of x must be (batch, time, feat)/(N, T, F) - if h_0 is None: # Create initial hidden states if necessary - h_0 = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device) + if h_0 is None and self.backbone_type != 'tres_deltagru': + # Create directly on the input device. TRes-DeltaGRU owns five + # recurrent states internally and historically discarded this one. + h_0 = torch.zeros( + self.num_layers, batch_size, self.hidden_size, device=x.device + ) # Forward Propagate through the RNN out = self.backbone(x, h_0) diff --git a/modules/cuda_graph_training.py b/modules/cuda_graph_training.py new file mode 100644 index 0000000..cdcb3c7 --- /dev/null +++ b/modules/cuda_graph_training.py @@ -0,0 +1,377 @@ +"""Opt-in CUDA-graph training for the supported DPD cascade. + +Only the existing clean-eager forward, MSE loss, gradient zeroing, and backward +are captured. Batch copies, clipping, and optimizer updates remain ordinary +PyTorch operations so parameters and optimizer state can change in place +between replays. Every unsupported condition returns ``None`` for a strict +clean-eager fallback. +""" + +from __future__ import annotations + +import contextlib +import os +import threading +import weakref +from dataclasses import dataclass + +import torch +from torch import nn + +from backbones.cuda_graph_frozen_dgru import disable_cuda_graph_frozen_dgru + + +_DISABLE_ENV = "OPENDPD_DISABLE_CUDA_GRAPH_TRAINING" +_TRUE_VALUES = {"1", "true", "yes", "on"} +_CACHES = weakref.WeakKeyDictionary() +_CACHES_LOCK = threading.Lock() +_CAPTURE_FAILED = object() + + +@dataclass(frozen=True) +class CudaGraphTrainingResult: + """Result of one captured forward/loss/backward replay.""" + + loss: torch.Tensor + output: torch.Tensor + + +def _disabled() -> bool: + return os.getenv(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES + + +def _cuda_autocast_enabled() -> bool: + try: + return torch.is_autocast_enabled("cuda") + except TypeError: # pragma: no cover - compatibility with older PyTorch. + return torch.is_autocast_enabled() + + +def _tres_layers(module: nn.Module) -> tuple[nn.Module, ...]: + """Return only actual TRes recurrence layers, not lookalike attributes.""" + + from backbones.tres_deltagru import DeltaGRULayer + + return tuple( + child for child in module.modules() + if isinstance(child, DeltaGRULayer) + ) + + +@contextlib.contextmanager +def force_clean_eager_tres(module: nn.Module): + """Temporarily disable fused TRes arithmetic, restoring it afterward.""" + + layers = _tres_layers(module) + previous = tuple(layer.use_triton for layer in layers) + for layer in layers: + layer.use_triton = False + try: + yield + finally: + for layer, enabled in zip(layers, previous, strict=True): + layer.use_triton = enabled + + +def _supported_cascade(module: nn.Module) -> bool: + """Limit capture to the measured, state-free DPD/PA architecture.""" + + from backbones.dgru import DGRU + from backbones.tres_deltagru import TResDeltaGRU + from models import CascadedModel, CoreModel + + if not isinstance(module, CascadedModel): + return False + dpd = module.dpd_model + pa = module.pa_model + if not isinstance(dpd, CoreModel) or not isinstance(pa, CoreModel): + return False + if ( + dpd.backbone_type != "tres_deltagru" + or pa.backbone_type != "dgru" + or not isinstance(dpd.backbone, TResDeltaGRU) + or not isinstance(pa.backbone, DGRU) + or dpd.backbone.rnn.debug + or any(parameter.requires_grad for parameter in pa.parameters()) + ): + return False + return tuple(_tres_layers(module)) == (dpd.backbone.rnn,) + + +def _can_capture( + module: nn.Module, + features: torch.Tensor, + targets: torch.Tensor, + criterion: nn.Module, + parameters: tuple[torch.Tensor, ...], +) -> bool: + if ( + _disabled() + or not module.training + or not torch.is_grad_enabled() + or not isinstance(criterion, nn.MSELoss) + or criterion.reduction != "mean" + or not _supported_cascade(module) + or not parameters + or features.requires_grad + or targets.requires_grad + or not features.is_cuda + or not targets.is_cuda + or features.device != targets.device + or features.dtype != torch.float32 + or targets.dtype != torch.float32 + or not features.is_contiguous() + or not targets.is_contiguous() + or torch.version.cuda is None + or getattr(torch.version, "hip", None) is not None + or _cuda_autocast_enabled() + or not hasattr(torch.cuda, "CUDAGraph") + ): + return False + try: + if torch.cuda.is_current_stream_capturing(): + return False + except RuntimeError: + return False + + model_parameters = tuple(module.parameters()) + if any( + child._forward_hooks + or child._forward_pre_hooks + or child._backward_hooks + or child._backward_pre_hooks + for child in module.modules() + ): + return False + trainable = tuple( + parameter for parameter in model_parameters + if parameter.requires_grad + ) + if parameters != trainable: + return False + return all( + parameter.device == features.device + and parameter.dtype == torch.float32 + for parameter in model_parameters + ) + + +def _cache_key( + module: nn.Module, + features: torch.Tensor, + targets: torch.Tensor, + parameters: tuple[torch.Tensor, ...], +): + return ( + features.device.index, + features.dtype, + tuple(features.shape), + tuple(features.stride()), + tuple(targets.shape), + tuple(targets.stride()), + tuple(parameter.data_ptr() for parameter in module.parameters()), + tuple(parameter.data_ptr() for parameter in parameters), + ) + + +class _GraphCache: + def __init__(self): + self.entries = {} + self.capture_streams = {} + self.lock = threading.Lock() + + def get_or_create( + self, module, features, targets, criterion, parameters + ): + key = _cache_key(module, features, targets, parameters) + entry = self.entries.get(key) + if entry is _CAPTURE_FAILED: + return None + if entry is not None and entry.failed: + return None + if entry is not None and entry.is_valid(): + return entry + + # A simultaneous first call falls back rather than attempting a second + # capture of the same module and parameter storage. + if not self.lock.acquire(blocking=False): + return None + try: + entry = self.entries.get(key) + if entry is _CAPTURE_FAILED: + return None + if entry is not None and entry.failed: + return None + if entry is None or not entry.is_valid(): + try: + device_index = features.device.index + capture_stream = self.capture_streams.get(device_index) + if capture_stream is None: + capture_stream = torch.cuda.Stream( + device=features.device + ) + self.capture_streams[device_index] = capture_stream + entry = _CapturedTrainingStep( + module, features, targets, criterion, parameters, + capture_stream, + ) + except (RuntimeError, torch.cuda.OutOfMemoryError): + self.entries[key] = _CAPTURE_FAILED + return None + self.entries[key] = entry + return entry + finally: + self.lock.release() + + +class _CapturedTrainingStep: + def __init__( + self, module, example_features, example_targets, + criterion, parameters, capture_stream + ): + self.parameters = parameters + self.parameter_pointers = tuple( + parameter.data_ptr() for parameter in parameters + ) + self.lock = threading.Lock() + self.busy = False + self.failed = False + self.pending_event = None + self.pending_stream = None + + with torch.cuda.device(example_features.device), torch.enable_grad(): + self.capture_stream = capture_stream + self.capture_stream.wait_stream(torch.cuda.current_stream()) + with ( + force_clean_eager_tres(module), + disable_cuda_graph_frozen_dgru(), + torch.cuda.stream(self.capture_stream), + ): + self.static_features = example_features.detach().clone() + self.static_targets = example_targets.detach().clone() + # Initialize cuDNN/cuBLAS and persistent parameter-grad buffers + # on the exact stream that will own the capture. + for _ in range(2): + for parameter in parameters: + if parameter.grad is not None: + parameter.grad.zero_() + warm_output = module(self.static_features) + warm_loss = criterion(warm_output, self.static_targets) + warm_loss.backward() + del warm_output, warm_loss + + torch.cuda.current_stream().wait_stream(self.capture_stream) + torch.cuda.synchronize(example_features.device) + if any(parameter.grad is None for parameter in parameters): + raise RuntimeError( + "captured parameters must all receive gradients" + ) + + self.graph = torch.cuda.CUDAGraph() + with ( + force_clean_eager_tres(module), + disable_cuda_graph_frozen_dgru(), + torch.cuda.graph(self.graph, stream=self.capture_stream), + ): + for parameter in parameters: + parameter.grad.zero_() + self.static_output = module(self.static_features) + self.static_loss = criterion( + self.static_output, self.static_targets + ) + self.static_loss.backward() + + self.gradient_pointers = tuple( + parameter.grad.data_ptr() for parameter in parameters + ) + + def is_valid(self) -> bool: + return ( + not self.failed + and self.parameter_pointers == tuple( + parameter.data_ptr() for parameter in self.parameters + ) + and all( + parameter.grad is not None for parameter in self.parameters + ) + and self.gradient_pointers == tuple( + parameter.grad.data_ptr() for parameter in self.parameters + ) + ) + + def reserve(self) -> bool: + with self.lock: + if not self.is_valid() or self.busy: + return False + if self.pending_event is not None: + stream = torch.cuda.current_stream( + self.static_features.device + ) + same_stream = stream.cuda_stream == self.pending_stream + if not same_stream and not self.pending_event.query(): + return False + self.pending_event = None + self.pending_stream = None + self.busy = True + return True + + def release(self, *, failed: bool = False) -> None: + with self.lock: + self.failed = self.failed or failed + stream = torch.cuda.current_stream(self.static_features.device) + event = torch.cuda.Event() + event.record(stream) + self.pending_event = event + self.pending_stream = stream.cuda_stream + self.busy = False + + def run(self, module, features, targets): + if not self.reserve(): + return None + try: + self.static_features.copy_(features) + self.static_targets.copy_(targets) + with force_clean_eager_tres(module): + self.graph.replay() + # Each logged loss needs independent storage because replay updates + # the static scalar on the next batch. + result = CudaGraphTrainingResult( + loss=self.static_loss.detach().clone(), + output=self.static_output.detach(), + ) + except (RuntimeError, torch.cuda.OutOfMemoryError): + self.release(failed=True) + return None + self.release() + return result + + +def try_cuda_graph_training_step( + module: nn.Module, + features: torch.Tensor, + targets: torch.Tensor, + criterion: nn.Module, + parameters: tuple[torch.Tensor, ...], +) -> CudaGraphTrainingResult | None: + """Replay one captured training body, or return ``None`` for fallback.""" + + if not _can_capture(module, features, targets, criterion, parameters): + return None + with _CACHES_LOCK: + cache = _CACHES.get(module) + if cache is None: + cache = _GraphCache() + _CACHES[module] = cache + step = cache.get_or_create( + module, features, targets, criterion, parameters + ) + if step is None: + return None + return step.run(module, features, targets) + + +def clear_cuda_graph_training_cache() -> None: + """Drop all captured training steps; intended primarily for tests.""" + + with _CACHES_LOCK: + _CACHES.clear() diff --git a/modules/data_collector.py b/modules/data_collector.py index e0ec4e5..d65f55e 100644 --- a/modules/data_collector.py +++ b/modules/data_collector.py @@ -233,18 +233,40 @@ def __getitem__(self, idx): class IQFrameDataset(Dataset): def __init__(self, features, targets, frame_length, stride=1): # Convert segments into frames - self.features = torch.Tensor(self.get_frames(features, frame_length, stride)) - self.targets = torch.Tensor(self.get_frames(targets, frame_length, stride)) + # ``torch.Tensor(data)`` historically converts every input dtype to + # float32. Keep that behavior while forcing independent storage. + self.features = torch.tensor( + self.get_frames(features, frame_length, stride), dtype=torch.float32 + ) + self.targets = torch.tensor( + self.get_frames(targets, frame_length, stride), dtype=torch.float32 + ) @staticmethod def get_frames(sequence, frame_length, stride_length): - frames = [] - sequence_length = len(sequence) - num_frames = (sequence_length - frame_length) // stride_length + 1 - for i in range(num_frames): - frame = sequence[i * stride_length: i * stride_length + frame_length] - frames.append(frame) - return np.stack(frames) + sequence = np.asarray(sequence) + if frame_length <= 0: + raise ValueError("frame_length must be positive") + if stride_length <= 0: + raise ValueError("stride_length must be positive") + if len(sequence) < frame_length: + raise ValueError("frame_length cannot exceed the sequence length") + + # sliding_window_view creates every overlapping frame at once. For a + # two-dimensional IQ sequence it initially places the window axis last + # (N, IQ, frame); move it behind the sample axis to recover the exact + # historical (N, frame, IQ) order, then apply the requested stride. + frames = np.lib.stride_tricks.sliding_window_view( + sequence, window_shape=frame_length, axis=0 + ) + if sequence.ndim > 1: + frames = np.moveaxis(frames, -1, 1) + frames = frames[::stride_length] + + # Window arrays alias the source by construction. Return a contiguous + # copy so neither the dataset nor callers observe storage aliasing. + return np.array(frames, copy=True, order='C') + def __len__(self): return len(self.features) diff --git a/modules/paths.py b/modules/paths.py index 5c45ccd..c35323f 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -4,6 +4,7 @@ import argparse import os +import warnings def gen_log_stat(args: argparse.Namespace, elapsed_time, net, optimizer, epoch, train_stat=None, val_stat=None, @@ -45,18 +46,30 @@ def gen_log_stat(args: argparse.Namespace, elapsed_time, net, optimizer, epoch, 'HIDDEN_SIZE': hidden_size, } - # Add threshold values to log - if args.step == 'train_dpd': - if 'delta' in net.dpd_model.backbone_type: - log_stat['THX'] = net.dpd_model.backbone.thx - log_stat['THH'] = net.dpd_model.backbone.thh + # Record delta thresholds for both PA and DPD training. They are runtime + # model settings rather than checkpoint tensors, so the CSV evidence is + # what lets a benchmark prove which thresholds were actually used. + delta_model = None + if args.step == 'train_pa' and 'delta' in net.backbone_type: + delta_model = net + elif args.step == 'train_dpd' and 'delta' in net.dpd_model.backbone_type: + delta_model = net.dpd_model + + if delta_model is not None: + log_stat['THX'] = delta_model.backbone.thx + log_stat['THH'] = delta_model.backbone.thh - # Add sparsity metrics if available - if 'delta' in net.dpd_model.backbone_type: - sparsity_metrics = net.dpd_model.backbone.get_temporal_sparsity() - sparsity_log = {f'{k}': v for k, v in sparsity_metrics.items()} - log_stat.update(sparsity_log) - net.dpd_model.backbone.set_debug(1) + # Statistics are opt-in because collecting them in an eager recurrent + # cell adds reductions at every timestep. + if ( + getattr(delta_model.backbone.rnn, 'debug', 0) + and hasattr(delta_model.backbone, 'get_temporal_sparsity') + and hasattr(delta_model.backbone, 'set_debug') + ): + sparsity_metrics = delta_model.backbone.get_temporal_sparsity() + sparsity_log = {f'{k}': v for k, v in sparsity_metrics.items()} + log_stat.update(sparsity_log) + delta_model.backbone.set_debug(1) # Merge stat dicts into the log dict if train_stat is not None: @@ -95,6 +108,33 @@ def gen_file_paths(path_dir_save: str, path_dir_log_hist: str, path_dir_log_best return file_paths +def warn_if_model_artifacts_exist(model_id: str, file_paths): + """Warn before an architecture-only model ID reuses existing artifacts. + + Model IDs are part of the public checkpoint lookup convention, so changing + their format would break existing consumers. They do not, however, encode + the full training recipe. An explicit warning keeps that convention while + making a potentially destructive recipe collision visible before training. + """ + existing_paths = [path for path in file_paths if os.path.exists(path)] + if not existing_paths: + return + + formatted_paths = ", ".join(os.path.normpath(path) for path in existing_paths) + warnings.warn( + ( + f"Existing artifact(s) for model ID '{model_id}' will be reused or " + f"overwritten: {formatted_paths}. This model ID does not encode the " + "full training recipe (for example optimizer, learning rate, " + "schedule, batch size, and epoch count), so the existing artifacts " + "may come from a different recipe. Move or rename them before " + "continuing if they must be preserved." + ), + RuntimeWarning, + stacklevel=2, + ) + + def create_folder(folder_list): for folder in folder_list: try: @@ -115,4 +155,4 @@ def gen_pa_model_id(args): list_pamodel_id += list(item) pa_model_id = '_'.join(list_pamodel_id) pa_model_id = 'PA_' + pa_model_id - return pa_model_id \ No newline at end of file + return pa_model_id diff --git a/modules/train_funcs.py b/modules/train_funcs.py index 1fa1419..5b0b3a4 100644 --- a/modules/train_funcs.py +++ b/modules/train_funcs.py @@ -2,6 +2,8 @@ __license__ = "Apache-2.0 License" __email__ = "yizhuo.wu@tudelft.nl, chang.gao@tudelft.nl" +import contextlib + import numpy as np import torch import torch.nn as nn @@ -13,41 +15,82 @@ import argparse +def _mean_losses(losses): + """Return the historical NumPy mean with one device-to-host transfer. + + ``Tensor.item()`` used to synchronize the accelerator once per batch. A + detached scalar keeps exactly the same stored value, so we can transfer all + batch losses together and still let NumPy perform the same float mean as + before. + """ + if not losses: + # Preserve NumPy's historical empty-loader result (NaN). + return np.mean([]) + return np.mean(torch.stack(losses).cpu().tolist()) + + def net_train(log: Dict[str, Any], net: nn.Module, dataloader: DataLoader, optimizer: Optimizer, criterion: Callable, grad_clip_val: float, - device: torch.device): + device: torch.device, + cuda_graph_training: bool = False): # Set Network to Training Mode net = net.train() # Statistics losses = [] + # The optimizer contains only trainable parameters in the standard + # pipeline. Materialize this tuple once per epoch rather than walking the + # complete cascaded model for every batch. + trainable_params = tuple( + parameter + for group in optimizer.param_groups + for parameter in group['params'] + if parameter.requires_grad + ) + non_blocking = device.type == 'cuda' + if cuda_graph_training: + from modules.cuda_graph_training import ( + force_clean_eager_tres, + try_cuda_graph_training_step, + ) # Iterate through batches for features, targets in tqdm(dataloader): # Move features and targets to the proper device - features = features.to(device) - targets = targets.to(device) - # Initialize all gradients to zero - optimizer.zero_grad() - # Forward Propagation - out = net(features) - # Calculate the Loss Function - loss = criterion(out, targets) - # Backward propagation - loss.backward() + features = features.to(device, non_blocking=non_blocking) + targets = targets.to(device, non_blocking=non_blocking) + graph_result = None + if cuda_graph_training: + graph_result = try_cuda_graph_training_step( + net, features, targets, criterion, trainable_params + ) + if graph_result is None: + # Strict opt-in fallback keeps the same clean-eager TRes + # arithmetic as the captured body. The original dispatch flag is + # restored immediately afterward, including for evaluation. + context = ( + force_clean_eager_tres(net) + if cuda_graph_training else contextlib.nullcontext() + ) + with context: + optimizer.zero_grad(set_to_none=True) + out = net(features) + loss = criterion(out, targets) + loss.backward() + else: + loss = graph_result.loss # Gradient clipping if grad_clip_val != 0: - nn.utils.clip_grad_norm_(net.parameters(), grad_clip_val) + nn.utils.clip_grad_norm_(trainable_params, grad_clip_val) # Update parameters optimizer.step() - # Detach loss from the graph indicating the end of forward propagation - loss.detach() - # Get losses - losses.append(loss.item()) + # Keep only the scalar value, not its graph. The complete vector is + # copied to the host once after the epoch to avoid a sync per batch. + losses.append(loss.detach()) # Average loss - loss = np.mean(losses) + loss = _mean_losses(losses) # Save Statistics log['loss'] = loss # End of Training Epoch @@ -60,30 +103,32 @@ def net_eval(log: Dict, criterion: Callable, device: torch.device): net = net.eval() - with torch.no_grad(): + with torch.inference_mode(): # Statistics losses = [] prediction = [] ground_truth = [] + non_blocking = device.type == 'cuda' # Batch Iteration for features, targets in tqdm(dataloader): # Move features and targets to the proper device - features = features.to(device) - targets = targets.to(device) + features = features.to(device, non_blocking=non_blocking) + targets = targets.to(device, non_blocking=non_blocking) # Forward Propagation outputs = net(features) # Calculate loss function loss = criterion(outputs, targets) # Collect prediction and ground truth for metric calculation - prediction.append(outputs.cpu()) - ground_truth.append(targets.cpu()) + prediction.append(outputs) + ground_truth.append(targets) # Collect losses to calculate the average loss per epoch - losses.append(loss.item()) + losses.append(loss) # Average loss per epoch - avg_loss = np.mean(losses) - # Prediction and Ground Truth - prediction = torch.cat(prediction, dim=0).numpy() - ground_truth = torch.cat(ground_truth, dim=0).numpy() + avg_loss = _mean_losses(losses) + # Concatenate before copying so evaluation performs one host transfer per + # result tensor instead of one transfer per batch. + prediction = torch.cat(prediction, dim=0).cpu().numpy() + ground_truth = torch.cat(ground_truth, dim=0).cpu().numpy() # Save Statistics log['loss'] = avg_loss # End of Evaluation Epoch @@ -92,7 +137,14 @@ def net_eval(log: Dict, def calculate_metrics(args: argparse.Namespace, stat: Dict[str, Any], prediction: np.ndarray, ground_truth: np.ndarray): stat['NMSE'] = metrics.NMSE(prediction, ground_truth) - stat['EVM'] = metrics.EVM(prediction, ground_truth, bw_main_ch=args.bw_main_ch, n_sub_ch=args.n_sub_ch, nperseg=args.nperseg) + stat['EVM'] = metrics.EVM( + prediction, + ground_truth, + sample_rate=args.input_signal_fs, + bw_main_ch=args.bw_main_ch, + n_sub_ch=args.n_sub_ch, + nperseg=args.nperseg, + ) ACLR_L = [] ACLR_R = [] ACLR_left, ACLR_right = metrics.ACLR(prediction, fs=args.input_signal_fs, nperseg=args.nperseg, diff --git a/opendpd/api.py b/opendpd/api.py index 43348c9..ffb608a 100644 --- a/opendpd/api.py +++ b/opendpd/api.py @@ -25,14 +25,29 @@ from arguments import get_arguments +def _append_cli_kwargs(kwargs: Dict[str, Any]) -> None: + """Append keyword options while respecting argparse boolean flags.""" + + boolean_flags = { + 'use_segments', 'quant', 'collect_delta_stats', + 'cuda_graph_training', 'plot', + } + for key, value in kwargs.items(): + if isinstance(value, bool) and key in boolean_flags: + if value: + sys.argv.append(f'--{key}') + elif value is not None: + sys.argv.extend([f'--{key}', str(value)]) + + def train_pa( dataset_name: Optional[str] = None, dataset_path: Optional[str] = None, PA_backbone: str = 'gru', PA_hidden_size: int = 23, - n_epochs: int = 100, - batch_size: int = 256, - lr: float = 5e-4, + n_epochs: int = 300, + batch_size: int = 64, + lr: float = 5e-3, accelerator: str = 'cpu', frame_length: int = 200, seed: int = 0, @@ -98,8 +113,7 @@ def train_pa( sys.argv.extend(['--plot_every', str(plot_every)]) # Add any additional keyword arguments - for key, value in kwargs.items(): - sys.argv.extend([f'--{key}', str(value)]) + _append_cli_kwargs(kwargs) # Create project and run training proj = Project() @@ -119,9 +133,9 @@ def train_dpd( DPD_hidden_size: int = 15, PA_backbone: str = 'gru', PA_hidden_size: int = 23, - n_epochs: int = 100, - batch_size: int = 256, - lr: float = 5e-4, + n_epochs: int = 300, + batch_size: int = 64, + lr: float = 5e-3, accelerator: str = 'cpu', frame_length: int = 200, seed: int = 0, @@ -129,6 +143,8 @@ def train_dpd( thh: float = 0.0, plot: bool = False, plot_every: int = 1, + collect_delta_stats: bool = False, + cuda_graph_training: bool = False, **kwargs ) -> Dict[str, Any]: """ @@ -151,6 +167,8 @@ def train_dpd( thh: Threshold for hidden state deltas (for delta-based models) plot: Enable plot generation during training plot_every: Generate per-epoch plots every N epochs + collect_delta_stats: Collect temporal sparsity diagnostics during training + cuda_graph_training: Opt in to guarded CUDA-graph DPD training **kwargs: Additional arguments passed to the training configuration Returns: @@ -188,13 +206,16 @@ def train_dpd( sys.argv.extend(['--seed', str(seed)]) sys.argv.extend(['--thx', str(thx)]) sys.argv.extend(['--thh', str(thh)]) + if collect_delta_stats: + sys.argv.append('--collect_delta_stats') + if cuda_graph_training: + sys.argv.append('--cuda_graph_training') if plot: sys.argv.append('--plot') sys.argv.extend(['--plot_every', str(plot_every)]) # Add any additional keyword arguments - for key, value in kwargs.items(): - sys.argv.extend([f'--{key}', str(value)]) + _append_cli_kwargs(kwargs) # Create project and run training proj = Project() @@ -258,8 +279,7 @@ def run_dpd( sys.argv.append('--plot') # Add any additional keyword arguments - for key, value in kwargs.items(): - sys.argv.extend([f'--{key}', str(value)]) + _append_cli_kwargs(kwargs) # Create project and run DPD proj = Project() @@ -318,8 +338,7 @@ def plot_dpd( sys.argv.extend(['--DPD_hidden_size', str(DPD_hidden_size)]) sys.argv.extend(['--accelerator', accelerator]) - for key, value in kwargs.items(): - sys.argv.extend([f'--{key}', str(value)]) + _append_cli_kwargs(kwargs) proj = Project() plot_module.main(proj) @@ -571,4 +590,3 @@ def run(self, **kwargs): config['dataset_path'] = self.dataset_path return run_dpd(**config) - diff --git a/project.py b/project.py index 9f2878c..717d301 100644 --- a/project.py +++ b/project.py @@ -13,7 +13,13 @@ from torch import optim from torch.utils.data import DataLoader from arguments import get_arguments -from modules.paths import create_folder, gen_log_stat, gen_dir_paths, gen_file_paths +from modules.paths import ( + create_folder, + gen_log_stat, + gen_dir_paths, + gen_file_paths, + warn_if_model_artifacts_exist, +) from modules.train_funcs import net_train, net_eval, calculate_metrics from utils import util from modules.loggers import PandasLogger @@ -94,6 +100,7 @@ def gen_dpd_model_id(self, n_net_params): def build_logger(self, model_id: str): # Get Save and Log Paths file_paths = gen_file_paths(self.path_dir_save, self.path_dir_log_hist, self.path_dir_log_best, model_id) + warn_if_model_artifacts_exist(model_id, file_paths) self.path_save_file_best, self.path_log_file_hist, self.path_log_file_best = file_paths print("::: Best Model Save Path: ", self.path_save_file_best) print("::: Log-History Path: ", self.path_log_file_hist) @@ -216,9 +223,19 @@ def build_dataloaders(self): test_set = IQSegmentDataset(X_test, y_test, nperseg=self.args.nperseg) # Define PyTorch Dataloaders - train_loader = DataLoader(train_set, batch_size=self.batch_size, shuffle=True) - val_loader = DataLoader(val_set, batch_size=self.batch_size_eval, shuffle=False) - test_loader = DataLoader(test_set, batch_size=self.batch_size_eval, shuffle=False) + pin_memory = self.device.type == 'cuda' + train_loader = DataLoader( + train_set, batch_size=self.batch_size, shuffle=True, + pin_memory=pin_memory + ) + val_loader = DataLoader( + val_set, batch_size=self.batch_size_eval, shuffle=False, + pin_memory=pin_memory + ) + test_loader = DataLoader( + test_set, batch_size=self.batch_size_eval, shuffle=False, + pin_memory=pin_memory + ) return (train_loader, val_loader, test_loader), input_size @@ -255,18 +272,33 @@ def build_criterion(self): raise AttributeError('Please use a valid loss function. Check argument.py.') def build_optimizer(self, net: nn.Module): + # Frozen PA parameters in a cascaded DPD model never receive gradients; + # excluding them avoids needless optimizer and clipping traversal while + # leaving the DPD updates unchanged. + trainable_params = tuple( + parameter for parameter in net.parameters() if parameter.requires_grad + ) + if not trainable_params: + raise ValueError("Cannot build an optimizer without trainable parameters") + # Optimizer if self.opt_type == 'adam': - optimizer = optim.Adam(net.parameters(), lr=self.lr) + optimizer = optim.Adam(trainable_params, lr=self.lr) elif self.opt_type == 'sgd': - optimizer = optim.SGD(net.parameters(), lr=self.lr, momentum=0.9) + optimizer = optim.SGD(trainable_params, lr=self.lr, momentum=0.9) elif self.opt_type == 'rmsprop': - optimizer = optim.RMSprop(net.parameters(), lr=self.lr) + optimizer = optim.RMSprop(trainable_params, lr=self.lr) elif self.opt_type == 'adamw': - optimizer = optim.AdamW(net.parameters(), lr=self.lr) + optimizer = optim.AdamW( + trainable_params, + lr=self.lr, + weight_decay=0.01, + betas=(0.9, 0.999), + eps=1e-8, + ) elif self.opt_type == 'adabound': import adabound # Run pip install adabound (https://github.com/Luolc/AdaBound) - optimizer = adabound.AdaBound(net.parameters(), lr=self.lr, final_lr=0.1) + optimizer = adabound.AdaBound(trainable_params, lr=self.lr, final_lr=0.1) else: raise RuntimeError('Please use a valid optimizer.') @@ -276,7 +308,10 @@ def build_optimizer(self, net: nn.Module): factor=self.decay_factor, patience=self.patience, threshold=1e-4, - min_lr=self.lr_end) + threshold_mode='rel', + cooldown=0, + min_lr=self.lr_end, + eps=1e-8) return optimizer, lr_scheduler @staticmethod @@ -352,7 +387,8 @@ def train(self, net: nn.Module, criterion: Callable, optimizer: optim.Optimizer, criterion=criterion, dataloader=train_loader, grad_clip_val=self.grad_clip_val, - device=self.device) + device=self.device, + cuda_graph_training=self.cuda_graph_training) # ----------- # Validation @@ -560,4 +596,4 @@ def train(self, net: nn.Module, criterion: Callable, optimizer: optim.Optimizer, print(f"Warning: GIF generation failed: {e}") print("Training Completed...") - print(" ") \ No newline at end of file + print(" ") diff --git a/steps/run_dpd.py b/steps/run_dpd.py index 8686f1c..45ee012 100644 --- a/steps/run_dpd.py +++ b/steps/run_dpd.py @@ -37,7 +37,9 @@ def main(proj: Project): hidden_size=proj.PA_hidden_size, num_layers=proj.PA_num_layers, backbone_type=proj.PA_backbone, - num_dvr_units=proj.num_dvr_units) + num_dvr_units=proj.num_dvr_units, + thx=proj.thx, + thh=proj.thh) n_net_pa_params = count_net_params(net_pa) print("::: Number of PA Model Parameters: ", n_net_pa_params) pa_model_id = proj.gen_pa_model_id(n_net_pa_params) diff --git a/steps/train_dpd.py b/steps/train_dpd.py index 3cfefd3..093cf68 100644 --- a/steps/train_dpd.py +++ b/steps/train_dpd.py @@ -30,7 +30,9 @@ def main(proj: Project): num_layers=proj.PA_num_layers, backbone_type=proj.PA_backbone, window_size=proj.window_size, - num_dvr_units=proj.num_dvr_units) + num_dvr_units=proj.num_dvr_units, + thx=proj.thx, + thh=proj.thh) n_net_pa_params = count_net_params(net_pa) print("::: Number of PA Model Parameters: ", n_net_pa_params) pa_model_id = proj.gen_pa_model_id(n_net_pa_params) @@ -50,6 +52,8 @@ def main(proj: Project): thh=proj.thh) net_dpd = get_quant_model(proj, net_dpd) + if proj.collect_delta_stats and hasattr(net_dpd.backbone, 'set_debug'): + net_dpd.backbone.set_debug(1) print("::: DPD Model: ", net_dpd) n_net_dpd_params = count_net_params(net_dpd) diff --git a/steps/train_pa.py b/steps/train_pa.py index be8a9db..0923440 100644 --- a/steps/train_pa.py +++ b/steps/train_pa.py @@ -26,7 +26,9 @@ def main(proj: Project): num_layers=proj.PA_num_layers, backbone_type=proj.PA_backbone, window_size=proj.window_size, - num_dvr_units=proj.num_dvr_units) + num_dvr_units=proj.num_dvr_units, + thx=proj.thx, + thh=proj.thh) n_net_pa_params = count_net_params(net) print("::: Number of PA Model Parameters: ", n_net_pa_params) pa_model_id = proj.gen_pa_model_id(n_net_pa_params) diff --git a/tests/test_api.py b/tests/test_api.py index 04e794f..b4a70e6 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,10 +1,12 @@ """Tests for the public opendpd Python API (opendpd/api.py).""" +import inspect import sys import numpy as np import pandas as pd import pytest +import torch import opendpd from conftest import REPO_ROOT, SMOKE_DATASET @@ -84,6 +86,92 @@ def test_rejects_wrong_columns(self, tmp_path): class TestApiTraining: + @pytest.mark.parametrize("function_name", ["train_pa", "train_dpd"]) + def test_benchmark_recipe_defaults_are_public_api_defaults(self, function_name): + signature = inspect.signature(getattr(opendpd, function_name)) + + assert signature.parameters["n_epochs"].default == 300 + assert signature.parameters["batch_size"].default == 64 + assert signature.parameters["lr"].default == 5e-3 + + def test_adamw_and_scheduler_recipe_is_explicit(self, monkeypatch): + import project as project_module + + calls = {} + optimizer_sentinel = object() + scheduler_sentinel = object() + + def capture_adamw(params, **kwargs): + calls["optimizer_params"] = tuple(params) + calls["optimizer_kwargs"] = kwargs + return optimizer_sentinel + + def capture_scheduler(**kwargs): + calls["scheduler_kwargs"] = kwargs + return scheduler_sentinel + + monkeypatch.setattr(project_module.optim, "AdamW", capture_adamw) + monkeypatch.setattr( + project_module.optim.lr_scheduler, + "ReduceLROnPlateau", + capture_scheduler, + ) + + project = project_module.Project.__new__(project_module.Project) + project.opt_type = "adamw" + project.lr = 5e-3 + project.decay_factor = 0.5 + project.patience = 5 + project.lr_end = 5e-5 + net = torch.nn.Linear(2, 2) + + optimizer, scheduler = project.build_optimizer(net) + + assert optimizer is optimizer_sentinel + assert scheduler is scheduler_sentinel + assert calls["optimizer_params"] == tuple(net.parameters()) + assert calls["optimizer_kwargs"] == { + "lr": 5e-3, + "weight_decay": 0.01, + "betas": (0.9, 0.999), + "eps": 1e-8, + } + assert calls["scheduler_kwargs"] == { + "optimizer": optimizer_sentinel, + "mode": "min", + "factor": 0.5, + "patience": 5, + "threshold": 1e-4, + "threshold_mode": "rel", + "cooldown": 0, + "min_lr": 5e-5, + "eps": 1e-8, + } + + def test_cuda_graph_training_flag_is_explicit_opt_in(self, preserved_argv): + from arguments import get_arguments + + sys.argv = ['opendpd'] + assert get_arguments().cuda_graph_training is False + sys.argv = ['opendpd', '--cuda_graph_training'] + assert get_arguments().cuda_graph_training is True + + def test_boolean_kwargs_are_emitted_as_flags(self, preserved_argv): + from opendpd import api + + sys.argv = ['opendpd'] + api._append_cli_kwargs({ + "collect_delta_stats": True, + "cuda_graph_training": True, + "plot": False, + "frame_stride": 16, + "unused": None, + }) + assert sys.argv == [ + 'opendpd', '--collect_delta_stats', '--cuda_graph_training', + '--frame_stride', '16' + ] + def test_train_pa_smoke(self, tmp_path, monkeypatch, preserved_argv): from pathlib import Path diff --git a/tests/test_benchmark_regressions.py b/tests/test_benchmark_regressions.py new file mode 100644 index 0000000..8cad8de --- /dev/null +++ b/tests/test_benchmark_regressions.py @@ -0,0 +1,985 @@ +"""Focused regression tests for the benchmark workflow.""" + +import copy +import importlib.util +import json +import os +from pathlib import Path +import shlex +import subprocess +import sys +import warnings + +import numpy as np +import pytest + +from project import Project + + +REPO_ROOT = Path(__file__).resolve().parent.parent +BENCHMARK_PATH = REPO_ROOT / "benchmark" / "benchmark_volterra.py" +BENCHMARK_SPEC = importlib.util.spec_from_file_location( + "benchmark_volterra", BENCHMARK_PATH +) +benchmark_volterra = importlib.util.module_from_spec(BENCHMARK_SPEC) +BENCHMARK_SPEC.loader.exec_module(benchmark_volterra) +COLLECTOR_PATH = REPO_ROOT / "benchmark" / "collect_benchmark_report.py" +COLLECTOR_SPEC = importlib.util.spec_from_file_location( + "collect_benchmark_report", COLLECTOR_PATH +) +collect_benchmark_report = importlib.util.module_from_spec(COLLECTOR_SPEC) +COLLECTOR_SPEC.loader.exec_module(collect_benchmark_report) + + +def parse_recorded_commands(path): + blocks = {} + lines = path.read_text().splitlines() + for index, line in enumerate(lines): + if line.startswith("# "): + label = line[2:] + assert label not in blocks, f"duplicate command label: {label}" + blocks[label] = shlex.split(lines[index + 1]) + return blocks + + +def option(command, name): + assert command.count(name) == 1, f"expected one {name} in {command}" + return command[command.index(name) + 1] + + +def polynomial_metrics(): + return { + "NMSE": -39.0, + "EVM": -47.0, + "ACLR_L": -42.0, + "ACLR_R": -40.0, + "ACLR_AVG": -41.0, + } + + +def polynomial_result(task="pa_modeling", model="mp"): + configuration = collect_benchmark_report.POLYNOMIAL_CONFIGURATIONS[task][model] + method = ( + "direct_least_squares" + if task == "pa_modeling" + else "indirect_learning_architecture" + ) + result = { + "schema_version": 4, + "dataset": "APA_200MHz", + "task": task, + "method": method, + "model": model, + "solver": "torch.linalg.lstsq", + "solver_mode": "gels", + "solver_driver": "gels", + "dtype": "torch.complex64", + "device": "cuda:0", + "solver_device": "cuda:0", + "solver_dtype": "torch.complex64", + "column_scaling": "l2", + "column_scale_min": 1.0, + "column_scale_max": 10.0, + "training_relative_residual": 0.1, + "coefficient_l2_norm": 1.0, + "regularization": None, + "svd_rcond": None, + "singular_value_max": None, + "singular_value_min": None, + "singular_value_cutoff": None, + "retained_singular_value_min": None, + "condition_number": None, + "full_rank_assumption": True, + "least_squares_rank": None, + "basis_configuration": configuration["basis_configuration"], + "complex_coefficients": configuration["complex_coefficients"], + "real_parameters": configuration["real_parameters"], + "parameter_count_convention": ( + "two real degrees of freedom per complex coefficient" + ), + "target_gain": 1.2, + "sample_rate_hz": 983_040_000.0, + "nperseg": 19_662, + "validation": polynomial_metrics(), + "test": polynomial_metrics(), + } + if task == "pa_modeling" and model == "gmp": + result.update( + { + "solver": "torch.linalg.svd", + "solver_mode": "truncated_svd", + "solver_driver": "gesvdj", + "full_rank_assumption": False, + "least_squares_rank": 650, + "regularization": "truncated_svd", + "svd_rcond": 1e-4, + "singular_value_max": 10.0, + "singular_value_min": 1e-6, + "singular_value_cutoff": 1e-3, + "retained_singular_value_min": 1.1e-3, + "condition_number": 1e7, + "coefficient_l2_norm": 8.0, + } + ) + if task == "dpd_ila": + result.update( + { + "pa_checkpoint": ( + "save/APA_200MHz/train_pa/" + "PA_S_0_M_TRES_GRU_H_27_F_200_P_2751.pt" + ), + "pa_checkpoint_sha256": "a" * 64, + "pa_evaluation_device": "cuda:0", + } + ) + return result + + +def test_volterra_cli_selects_distinct_mp_and_gmp_bases(): + parser = benchmark_volterra.build_arg_parser() + mp_args = parser.parse_args(["--dataset-name", "fixture", "--model", "mp"]) + gmp_args = parser.parse_args(["--dataset-name", "fixture", "--model", "gmp"]) + + mp_builder, mp_coefficients, mp_banner = benchmark_volterra.select_basis(mp_args) + gmp_builder, gmp_coefficients, gmp_banner = benchmark_volterra.select_basis( + gmp_args + ) + + signal = np.linspace(0.1, 1.0, 64) * np.exp( + 1j * np.linspace(0.0, 2.0 * np.pi, 64) + ) + mp_basis = mp_builder(signal) + gmp_basis = gmp_builder(signal) + + assert mp_coefficients == mp_args.K * mp_args.Q == 250 + assert gmp_coefficients == ( + gmp_args.Ka * gmp_args.La + + gmp_args.Kb * gmp_args.Lb * gmp_args.Mb + + gmp_args.Kc * gmp_args.Lc * gmp_args.Mc + ) == 255 + assert mp_basis.shape == (signal.size, mp_coefficients) + assert gmp_basis.shape == (signal.size, gmp_coefficients) + assert not np.allclose(mp_basis, gmp_basis[:, :mp_coefficients]) + assert " MP " in mp_banner + assert " GMP " in gmp_banner + + +def test_segmented_polynomial_basis_resets_delay_state(): + signal = np.arange(1, 9, dtype=np.float64).astype(np.complex128) + builder = lambda values: benchmark_volterra.build_mp_basis(values, K=1, Q=2) + + full = builder(signal) + segmented = benchmark_volterra.build_segmented_numpy_basis( + signal, builder, segment_length=4 + ) + + assert full[4, 1] == signal[3] + assert segmented[4, 1] == 0 + assert np.array_equal(segmented[:4], full[:4]) + + +def test_truncated_svd_stabilizes_rank_deficient_gmp_basis(): + import torch + + signal = np.ones(64, dtype=np.complex128) + target = 1.5 * signal + configuration = { + "Ka": 2, + "La": 2, + "Kb": 1, + "Lb": 2, + "Mb": 1, + "Kc": 1, + "Lc": 2, + "Mc": 1, + } + + coefficients, diagnostics = ( + benchmark_volterra.fit_complex_least_squares( + signal, + target, + polynomial_model="gmp", + configuration=configuration, + segment_length=64, + device=torch.device("cpu"), + dtype=torch.complex128, + solver_mode="truncated_svd", + svd_rcond=1e-4, + ) + ) + prediction = benchmark_volterra.apply_polynomial( + signal, + coefficients, + polynomial_model="gmp", + configuration=configuration, + segment_length=64, + ) + + assert diagnostics["solver_mode"] == "truncated_svd" + assert diagnostics["solver_implementation"] == "torch.linalg.svd" + assert diagnostics["regularization"] == "truncated_svd" + assert 0 < diagnostics["least_squares_rank"] < diagnostics["columns"] + assert diagnostics["svd_rcond"] == 1e-4 + assert np.all(np.isfinite(coefficients.numpy())) + assert np.linalg.norm(prediction - target) / np.linalg.norm(target) < 1e-12 + + +def test_solver_mode_rejects_incompatible_svd_cutoff(): + import torch + + signal = np.ones(16, dtype=np.complex128) + with pytest.raises(ValueError, match="only valid for truncated_svd"): + benchmark_volterra.fit_complex_least_squares( + signal, + signal, + polynomial_model="mp", + configuration={"K": 1, "Q": 1}, + segment_length=16, + device=torch.device("cpu"), + dtype=torch.complex128, + solver_mode="gels", + svd_rcond=1e-4, + ) + + +def test_build_logger_warns_before_existing_recipe_artifacts_are_overwritten( + tmp_path, +): + project = Project.__new__(Project) + project.path_dir_save = str(tmp_path / "save") + project.path_dir_log_hist = str(tmp_path / "history") + project.path_dir_log_best = str(tmp_path / "best") + project.log_precision = 8 + for directory in ( + project.path_dir_save, + project.path_dir_log_hist, + project.path_dir_log_best, + ): + Path(directory).mkdir() + + model_id = "PA_S_0_M_GRU_H_28_F_200_P_2746" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + project.build_logger(model_id) + assert caught == [] + + Path(project.path_save_file_best).touch() + with pytest.warns(RuntimeWarning, match="does not encode the full training recipe"): + project.build_logger(model_id) + + +def test_benchmark_report_runner_dry_run_records_complete_recipe(tmp_path): + output_dir = tmp_path / "benchmark-run" + runner = REPO_ROOT / "benchmark" / "reproduce_benchmark_report.sh" + environment = dict(os.environ, OPENDPD_PYTHON=sys.executable) + + result = subprocess.run( + [ + "bash", + str(runner), + "--dry-run", + "--output-dir", + str(output_dir), + ], + cwd=REPO_ROOT, + env=environment, + capture_output=True, + text=True, + timeout=30, + ) + + assert result.returncode == 0, result.stderr + commands = parse_recorded_commands(output_dir / "commands.log") + expected_labels = {"snapshot_context", "collect_report"} + expected_labels.update( + f"train_pa_{slug}_{model}" + for slug in ("apa", "dpa") + for model in ("gru", "tres_gru", "tres_deltagru") + ) + expected_labels.update( + f"train_dpd_{slug}_{experiment}" + for slug in ("apa", "dpa") + for experiment in ( + "gru_via_tres_gru_pa", + "tres_gru_via_tres_gru_pa", + "tres_deltagru_via_tres_gru_pa", + "tres_deltagru_via_tres_deltagru_pa", + ) + ) + expected_labels.update( + f"pa_model_{slug}_{model}" + for slug in ("apa", "dpa") + for model in ("mp", "gmp") + ) + expected_labels.update( + f"dpd_ila_{slug}_{model}" + for slug in ("apa", "dpa") + for model in ("mp", "gmp") + ) + assert set(commands) == expected_labels + + for slug, dataset in (("apa", "APA_200MHz"), ("dpa", "DPA_160MHz")): + for backbone, hidden_size in { + "gru": "28", + "tres_gru": "27", + "tres_deltagru": "27", + }.items(): + command = commands[f"train_pa_{slug}_{backbone}"] + assert command[1] == "main.py" + assert option(command, "--dataset_name") == dataset + assert option(command, "--step") == "train_pa" + assert option(command, "--frame_length") == "200" + assert option(command, "--frame_stride") == "1" + assert option(command, "--n_epochs") == "300" + assert option(command, "--opt_type") == "adamw" + assert option(command, "--batch_size") == "64" + assert option(command, "--batch_size_eval") == "64" + assert option(command, "--lr") == "5e-3" + assert option(command, "--lr_schedule") == "1" + assert option(command, "--lr_end") == "5e-5" + assert option(command, "--decay_factor") == "0.5" + assert option(command, "--patience") == "5" + assert option(command, "--loss_type") == "l2" + assert option(command, "--grad_clip_val") == "200" + assert option(command, "--thx") == "0" + assert option(command, "--thh") == "0" + assert option(command, "--seed") == "0" + assert option(command, "--re_level") == "soft" + assert option(command, "--accelerator") == "cuda" + assert option(command, "--devices") == "0" + assert option(command, "--PA_backbone") == backbone + assert option(command, "--PA_hidden_size") == hidden_size + + dpd_experiments = { + "gru_via_tres_gru_pa": ("gru", "16", "tres_gru"), + "tres_gru_via_tres_gru_pa": ( + "tres_gru", + "15", + "tres_gru", + ), + "tres_deltagru_via_tres_gru_pa": ( + "tres_deltagru", + "15", + "tres_gru", + ), + "tres_deltagru_via_tres_deltagru_pa": ( + "tres_deltagru", + "15", + "tres_deltagru", + ), + } + for experiment, (backbone, hidden_size, pa_backbone) in ( + dpd_experiments.items() + ): + command = commands[f"train_dpd_{slug}_{experiment}"] + assert command[1] == "main.py" + assert option(command, "--dataset_name") == dataset + assert option(command, "--step") == "train_dpd" + assert option(command, "--frame_length") == "200" + assert option(command, "--frame_stride") == "1" + assert option(command, "--n_epochs") == "300" + assert option(command, "--opt_type") == "adamw" + assert option(command, "--batch_size") == "64" + assert option(command, "--batch_size_eval") == "64" + assert option(command, "--lr") == "5e-3" + assert option(command, "--lr_schedule") == "1" + assert option(command, "--lr_end") == "5e-5" + assert option(command, "--decay_factor") == "0.5" + assert option(command, "--patience") == "5" + assert option(command, "--loss_type") == "l2" + assert option(command, "--grad_clip_val") == "200" + assert option(command, "--thx") == "0" + assert option(command, "--thh") == "0" + assert option(command, "--seed") == "0" + assert option(command, "--re_level") == "soft" + assert option(command, "--accelerator") == "cuda" + assert option(command, "--devices") == "0" + assert option(command, "--PA_backbone") == pa_backbone + assert option(command, "--PA_hidden_size") == "27" + assert option(command, "--DPD_backbone") == backbone + assert option(command, "--DPD_hidden_size") == hidden_size + + for slug, dataset in (("apa", "APA_200MHz"), ("dpa", "DPA_160MHz")): + pa_mp = commands[f"pa_model_{slug}_mp"] + pa_gmp = commands[f"pa_model_{slug}_gmp"] + dpd_mp = commands[f"dpd_ila_{slug}_mp"] + dpd_gmp = commands[f"dpd_ila_{slug}_gmp"] + for command, task in ( + (pa_mp, "pa_modeling"), + (pa_gmp, "pa_modeling"), + (dpd_mp, "dpd_ila"), + (dpd_gmp, "dpd_ila"), + ): + assert command[1:3] == ["-m", "benchmark.benchmark_volterra"] + assert option(command, "--task") == task + assert option(command, "--dataset-name") == dataset + assert option(command, "--solver-device") == "cuda:0" + assert option(command, "--solver-dtype") == "complex64" + assert option(pa_mp, "--solver-mode") == "gels" + assert option(pa_gmp, "--solver-mode") == "truncated_svd" + assert option(pa_gmp, "--svd-rcond") == "1e-4" + assert option(dpd_mp, "--solver-mode") == "gels" + assert option(dpd_gmp, "--solver-mode") == "gels" + assert "--svd-rcond" not in pa_mp + assert "--svd-rcond" not in dpd_mp + assert "--svd-rcond" not in dpd_gmp + assert option(pa_mp, "--K") == "9" + assert option(pa_mp, "--Q") == "150" + assert option(pa_gmp, "--La") == "30" + assert option(pa_gmp, "--Mb") == "5" + assert option(pa_gmp, "--Mc") == "5" + assert option(dpd_mp, "--K") == "5" + assert option(dpd_mp, "--Q") == "100" + assert option(dpd_gmp, "--La") == "20" + assert option(dpd_gmp, "--Mb") == "3" + assert option(dpd_gmp, "--Mc") == "2" + for command in (dpd_mp, dpd_gmp): + assert option(command, "--pa-backbone") == "tres_gru" + assert option(command, "--pa-hidden-size") == "27" + + assert "--snapshot-context" in commands["snapshot_context"] + assert "--snapshot-context" not in commands["collect_report"] + collect_benchmark_report.validate_neural_commands( + output_dir / "commands.log" + ) + collect_benchmark_report.validate_runner_commands( + output_dir / "commands.log", + output_dir=output_dir, + device=0, + ) + + recipe = json.loads((output_dir / "recipe.json").read_text()) + assert recipe == collect_benchmark_report.expected_runner_recipe() + assert recipe["schema_version"] == 6 + assert recipe["neural"]["epochs"] == 300 + assert recipe["neural"]["batch_size"] == 64 + assert recipe["neural"]["initial_learning_rate"] == 5e-3 + assert recipe["neural"]["scheduler_patience"] == 5 + assert recipe["neural"]["minimum_learning_rate"] == 5e-5 + assert recipe["neural"]["optimizer_weight_decay"] == 0.01 + assert recipe["neural"]["optimizer_betas"] == [0.9, 0.999] + assert recipe["neural"]["scheduler_threshold_mode"] == "rel" + assert recipe["pa_models"]["mp"]["real_parameters"] == 2700 + assert recipe["pa_models"]["tres_deltagru"] == { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + } + assert recipe["dpd_models"]["gmp"]["real_parameters"] == 1000 + assert recipe["polynomial_solvers"]["gmp_pa"] == { + "implementation": "torch.linalg.svd", + "mode": "truncated_svd", + "driver": "gesvdj", + "dtype": "torch.complex64", + "column_scaling": "l2", + "regularization": { + "type": "truncated_svd", + "relative_cutoff": 1e-4, + }, + } + assert recipe["dpd_tres_deltagru_pa_surrogate_comparison"] == { + "model": { + "hidden_size": 15, + "real_parameters": 999, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + }, + "pa_surrogates": { + "tres_gru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": None, + }, + "tres_deltagru": { + "hidden_size": 27, + "real_parameters": 2751, + "delta_thresholds": {"input": 0.0, "hidden": 0.0}, + }, + }, + } + + +def test_distributed_signal_papr_is_stable(): + apa_record = collect_benchmark_report.signal_record("APA_200MHz") + dpa_record = collect_benchmark_report.signal_record("DPA_160MHz") + apa = apa_record["papr"] + dpa = dpa_record["papr"] + + assert apa_record["samples"] == 98_304 + assert dpa_record["samples"] == 491_520 + assert apa["ccdf_probability"] == 1e-5 + assert apa["quantile_method"] == "numpy linear" + assert apa["ccdf_db"] == pytest.approx(10.00798, abs=1e-5) + assert apa["absolute_peak_db"] == pytest.approx(10.01054, abs=1e-5) + assert dpa["ccdf_db"] == pytest.approx(10.13155, abs=1e-5) + assert dpa["absolute_peak_db"] == pytest.approx(10.13976, abs=1e-5) + + +def test_pa_checkpoint_bindings_require_eight_unique_rows(tmp_path): + experiments = ( + ("gru", "tres_gru"), + ("tres_gru", "tres_gru"), + ("tres_deltagru", "tres_gru"), + ("tres_deltagru", "tres_deltagru"), + ) + rows = [ + ( + dataset, + dpd_model, + pa_model, + f"save/{dataset}/{pa_model}.pt", + digest * 64, + ) + for dataset, digest in (("APA_200MHz", "a"), ("DPA_160MHz", "b")) + for dpd_model, pa_model in experiments + ] + path = tmp_path / "pa_checkpoint_bindings.tsv" + path.write_text( + "dataset\tdpd_model_key\tpa_model_key\tcheckpoint\tsha256\n" + + "".join("\t".join(row) + "\n" for row in rows) + ) + + bindings = collect_benchmark_report.read_pa_bindings(tmp_path) + assert set(bindings) == { + (dataset, dpd_model, pa_model) + for dataset in ("APA_200MHz", "DPA_160MHz") + for dpd_model, pa_model in experiments + } + + rows[-1] = rows[0] + path.write_text( + "dataset\tdpd_model_key\tpa_model_key\tcheckpoint\tsha256\n" + + "".join("\t".join(row) + "\n" for row in rows) + ) + with pytest.raises(ValueError, match="Duplicate PA checkpoint binding"): + collect_benchmark_report.read_pa_bindings(tmp_path) + + +def test_polynomial_metric_validation_rejects_nonfinite_and_bad_average(): + metrics = polynomial_metrics() + assert collect_benchmark_report.normalize_polynomial_metrics(metrics) == { + "nmse_db": -39.0, + "evm_db": -47.0, + "aclr_left_db": -42.0, + "aclr_right_db": -40.0, + "aclr_avg_db": -41.0, + } + + with pytest.raises(ValueError, match="not finite"): + collect_benchmark_report.normalize_polynomial_metrics( + dict(metrics, NMSE=float("nan")) + ) + with pytest.raises(ValueError, match="ACLR average"): + collect_benchmark_report.normalize_polynomial_metrics( + dict(metrics, ACLR_AVG=-40.5) + ) + + +def benchmark_context(): + return { + "git": {"commit": "abc", "branch": "benchmark-fix", "dirty": False}, + "environment": {"python": "3.13"}, + "source_file_sha256": {"main.py": "1" * 64}, + "datasets": {"APA_200MHz": {"spec_sha256": "2" * 64}}, + } + + +def test_context_stability_accepts_identical_context(): + initial = benchmark_context() + collect_benchmark_report.validate_context_stability( + initial, copy.deepcopy(initial) + ) + + +@pytest.mark.parametrize( + ("path", "replacement", "message"), + [ + (("git", "commit"), "def", "Git commit"), + (("git", "branch"), "other", "Git branch"), + (("environment", "python"), "3.14", "software and hardware"), + (("source_file_sha256", "main.py"), "3" * 64, "source files"), + ( + ("datasets", "APA_200MHz", "spec_sha256"), + "4" * 64, + "dataset files", + ), + ], +) +def test_context_stability_rejects_drift(path, replacement, message): + initial = benchmark_context() + changed = copy.deepcopy(initial) + target = changed + for key in path[:-1]: + target = target[key] + target[path[-1]] = replacement + + with pytest.raises(ValueError, match=message): + collect_benchmark_report.validate_context_stability(initial, changed) + + +def test_context_loader_requires_pre_run_snapshot(tmp_path): + with pytest.raises(FileNotFoundError, match="Pre-run context snapshot"): + collect_benchmark_report.load_and_validate_context(tmp_path, 0) + + +def test_source_snapshot_covers_runtime_sources(): + sources = set(collect_benchmark_report.source_files()) + assert { + "arguments.py", + "models.py", + "project.py", + "backbones/gru.py", + "backbones/tres_deltagru.py", + "backbones/tres_gru.py", + "benchmark/benchmark_volterra.py", + "benchmark/collect_benchmark_report.py", + "modules/train_funcs.py", + "steps/train_dpd.py", + "steps/train_pa.py", + "utils/metrics.py", + } <= sources + + +def test_polynomial_pa_artifact_contract_rejects_recipe_drift(tmp_path): + canonical = polynomial_result(task="pa_modeling", model="mp") + output_path = tmp_path / "benchmark_report_apa_pa_mp.json" + output_path.write_text(json.dumps(canonical)) + + result = collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="pa_modeling", + model_name="mp", + ) + assert result["real_parameters"] == 2700 + assert result["least_squares_rank"] is None + assert result["pa_checkpoint_binding"] is None + + mutations = [ + ("solver", "other", "polynomial solver"), + ("basis_configuration", {"K": 9, "Q": 149}, "polynomial basis"), + ("complex_coefficients", 1349, "complex coefficient count"), + ("real_parameters", 2698, "parameter count"), + ("sample_rate_hz", 800_000_000.0, "sample rate"), + ("nperseg", 16_384, "segment length"), + ("column_scale_min", 0.0, "must be positive"), + ] + for field, replacement, message in mutations: + changed = copy.deepcopy(canonical) + changed[field] = replacement + output_path.write_text(json.dumps(changed)) + with pytest.raises(ValueError, match=message): + collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="pa_modeling", + model_name="mp", + ) + + +def test_gmp_pa_artifact_contract_requires_truncated_svd_evidence(tmp_path): + canonical = polynomial_result(task="pa_modeling", model="gmp") + output_path = tmp_path / "benchmark_report_apa_pa_gmp.json" + output_path.write_text(json.dumps(canonical)) + + result = collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="pa_modeling", + model_name="gmp", + ) + assert result["solver"] == "torch.linalg.svd" + assert result["solver_mode"] == "truncated_svd" + assert result["svd_rcond"] == 1e-4 + assert result["least_squares_rank"] == 650 + assert result["coefficient_l2_norm"] == 8.0 + + mutations = [ + ("solver_mode", "gels", "solver mode"), + ("svd_rcond", 1e-6, "SVD relative cutoff"), + ("least_squares_rank", 1350, "rank is outside"), + ("retained_singular_value_min", 5e-4, "does not exceed"), + ("condition_number", 1e3, "inconsistent"), + ] + for field, replacement, message in mutations: + changed = copy.deepcopy(canonical) + changed[field] = replacement + output_path.write_text(json.dumps(changed)) + with pytest.raises(ValueError, match=message): + collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="pa_modeling", + model_name="gmp", + ) + + +def test_polynomial_dpd_artifact_requires_reference_pa_hash(tmp_path): + canonical = polynomial_result(task="dpd_ila", model="mp") + output_path = tmp_path / "benchmark_report_apa_dpd_mp.json" + output_path.write_text(json.dumps(canonical)) + + result = collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="dpd_ila", + model_name="mp", + expected_pa_hash="a" * 64, + ) + assert result["real_parameters"] == 1000 + assert result["pa_checkpoint_binding"]["sha256"] == "a" * 64 + + changed = dict(canonical, pa_checkpoint_sha256="b" * 64) + output_path.write_text(json.dumps(changed)) + with pytest.raises(ValueError, match="checkpoint hash"): + collect_benchmark_report.collect_polynomial_model( + tmp_path, + dataset="APA_200MHz", + slug="apa", + task="dpd_ila", + model_name="mp", + expected_pa_hash="a" * 64, + ) + + +def test_checkpoint_validation_checks_loadability_and_parameter_count(tmp_path): + import torch + + checkpoint = tmp_path / "model.pt" + torch.save({"weight": torch.ones(3), "bias": torch.zeros(2)}, checkpoint) + collect_benchmark_report.validate_checkpoint(checkpoint, 5) + + with pytest.raises(ValueError, match="checkpoint parameter count"): + collect_benchmark_report.validate_checkpoint(checkpoint, 4) + + checkpoint.write_bytes(b"not a checkpoint") + with pytest.raises(ValueError, match="cannot be loaded safely"): + collect_benchmark_report.validate_checkpoint(checkpoint, 5) + + +def test_run_manifest_uses_v6_schema(): + assert collect_benchmark_report.RUN_SCHEMA_VERSION == 6 + + +def report_manifest(): + values = { + "APA_200MHz": { + "pa": [-37.0, -38.7, -38.8, -39.1, -39.2], + "dpd": [-45.2, -43.6, -51.4, -53.4], + "delta_dpd": [-54.1, -55.2], + }, + "DPA_160MHz": { + "pa": [-38.4, -39.4, -39.5, -39.7, -39.8], + "dpd": [-50.7, -52.8, -56.5, -59.5], + "delta_dpd": [-60.1, -61.3], + }, + } + display_names = { + "mp": "MP", + "gmp": "GMP", + "gru": "GRU-H28", + "tres_gru": "TRes-GRU-H27", + "tres_deltagru": "TRes-DeltaGRU-H27 (THX=THH=0)", + } + datasets = {} + for dataset, dataset_values in values.items(): + pa_models = {} + dpd_models = {} + delta_dpd_by_pa_surrogate = {} + for index, key in enumerate(collect_benchmark_report.PA_MODEL_ORDER): + polynomial = key in {"mp", "gmp"} + pa_models[key] = { + "method": ( + "direct_least_squares" + if polynomial + else "supervised_training" + ), + "model": { + "display_name": display_names[key], + "parameters": 2700 + index, + }, + "selected_epoch_zero_based": None if polynomial else 10 + index, + "metrics": { + "validation": { + "nmse_db": dataset_values["pa"][index] + 0.1, + "evm_db": -40.0 - index, + "aclr_left_db": -30.0 - index, + "aclr_right_db": -31.0 - index, + "aclr_avg_db": -30.5 - index, + }, + "test": { + "nmse_db": dataset_values["pa"][index], + "evm_db": -40.2 - index, + "aclr_left_db": -30.2 - index, + "aclr_right_db": -31.2 - index, + "aclr_avg_db": -30.7 - index, + }, + }, + } + if key == "gmp": + pa_models[key]["training_relative_residual"] = 0.0112 + pa_models[key]["regularization"] = "truncated_svd" + pa_models[key]["svd_rcond"] = 1e-4 + pa_models[key]["least_squares_rank"] = 650 + pa_models[key]["complex_coefficients"] = 1350 + if key == "tres_deltagru": + pa_models[key]["delta_thresholds"] = { + "input": 0.0, + "hidden": 0.0, + } + + for index, key in enumerate(collect_benchmark_report.DPD_MODEL_ORDER): + polynomial = key in {"mp", "gmp"} + dpd_models[key] = { + "method": ( + "indirect_learning_architecture" + if polynomial + else "direct_learning_architecture" + ), + "model": { + "display_name": display_names[key], + "parameters": 1000 + index, + }, + "selected_epoch_zero_based": None if polynomial else 20 + index, + "metrics": { + "validation": { + "nmse_db": -42.0 - index, + "evm_db": -45.0 - index, + "aclr_left_db": dataset_values["dpd"][index] - 0.4, + "aclr_right_db": dataset_values["dpd"][index] + 0.2, + "aclr_avg_db": dataset_values["dpd"][index] - 0.1, + }, + "test": { + "nmse_db": -42.2 - index, + "evm_db": -45.2 - index, + "aclr_left_db": dataset_values["dpd"][index] - 0.3, + "aclr_right_db": dataset_values["dpd"][index] + 0.3, + "aclr_avg_db": dataset_values["dpd"][index], + }, + }, + } + for index, pa_model_key in enumerate(("tres_gru", "tres_deltagru")): + aclr = dataset_values["delta_dpd"][index] + delta_dpd_by_pa_surrogate[pa_model_key] = { + "method": "direct_learning_architecture", + "model": { + "display_name": "TRes-DeltaGRU-H15 (THX=THH=0)", + "parameters": 999, + }, + "selected_epoch_zero_based": 30 + index, + "metrics": { + "validation": { + "nmse_db": -48.0 - index, + "evm_db": -50.0 - index, + "aclr_left_db": aclr - 0.3, + "aclr_right_db": aclr + 0.1, + "aclr_avg_db": aclr - 0.1, + }, + "test": { + "nmse_db": -48.2 - index, + "evm_db": -50.2 - index, + "aclr_left_db": aclr - 0.2, + "aclr_right_db": aclr + 0.2, + "aclr_avg_db": aclr, + }, + }, + "pa_surrogate": {"model_key": pa_model_key}, + } + datasets[dataset] = { + "pa_models": pa_models, + "dpd_models": dpd_models, + "dpd_tres_deltagru_by_pa_surrogate": ( + delta_dpd_by_pa_surrogate + ), + } + return { + "generated_at_utc": "2026-07-24T00:00:00+00:00", + "git": {"commit": "abc", "branch": "benchmark-fix", "dirty": False}, + "environment": { + "python": "3.13", + "torch": "2.0", + "cuda": "13.0", + "gpu": "Fixture GPU", + }, + "datasets": datasets, + } + + +def test_report_is_latest_only_and_documents_gmp_stability_and_temporal_context(): + markdown = collect_benchmark_report.render_markdown( + report_manifest(), + canonical=True, + ) + + assert "benchmark_results.png" in markdown + assert "benchmark_delta_dpd_results.png" in markdown + assert "APA GMP stability" in markdown + assert "rcond=1e-04" in markdown + assert "650/1,350 singular directions" in markdown + assert "Truncated SVD (rank 650/1,350)" in markdown + assert "16-sample right context" in markdown + assert "wraps the final position" in markdown + assert "not demodulated constellation EVM" in markdown + assert "reproduce_benchmark_report.sh" in markdown + assert "commands.log" not in markdown + assert "previous report" not in markdown.lower() + assert "DGRU" not in markdown + assert "TRes-DeltaGRU" in markdown + assert "THX=THH=0" in markdown + assert "dense zero-threshold recurrence" in markdown + assert "TRes-DeltaGRU DPD by PA surrogate" in markdown + assert "separately trained" in markdown + for dataset in ("APA_200MHz", "DPA_160MHz"): + dataset_section = markdown.split(f"## {dataset}", 1)[1].split( + "\n## ", + 1, + )[0] + primary_dpd = dataset_section.split("### DPD\n", 1)[1].split( + "### TRes-DeltaGRU DPD by PA surrogate", + 1, + )[0] + assert "TRes-DeltaGRU" not in primary_dpd + + +def test_results_figure_is_a_nonempty_png(tmp_path): + output = tmp_path / "benchmark_results.png" + + collect_benchmark_report.write_results_figure(report_manifest(), output) + + assert output.read_bytes().startswith(b"\x89PNG\r\n\x1a\n") + assert output.stat().st_size > 10_000 + + +def test_delta_dpd_figure_is_a_nonempty_png(tmp_path): + output = tmp_path / "benchmark_delta_dpd_results.png" + + collect_benchmark_report.write_delta_dpd_figure( + report_manifest(), + output, + ) + + assert output.read_bytes().startswith(b"\x89PNG\r\n\x1a\n") + assert output.stat().st_size > 10_000 + + +def test_source_snapshot_roundtrip_matches_pre_run_hashes(tmp_path): + paths = ["arguments.py", "benchmark/collect_benchmark_report.py"] + expected = { + path: collect_benchmark_report.sha256_file(REPO_ROOT / path) + for path in paths + } + output = tmp_path / "source_snapshot.tar.gz" + + collect_benchmark_report.write_source_snapshot(paths, output) + metadata = collect_benchmark_report.validate_source_snapshot( + output, + expected, + ) + + assert metadata["path"] == "source_snapshot.tar.gz" + assert metadata["file_count"] == len(paths) + assert metadata["sha256"] == collect_benchmark_report.sha256_file(output) diff --git a/tests/test_cli.py b/tests/test_cli.py index 30b6e10..984cd74 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,6 +4,7 @@ import subprocess import sys +from arguments import get_arguments from conftest import MAIN_PY, REPO_ROOT @@ -38,3 +39,23 @@ def test_main_py_rejects_unknown_backbone(): [sys.executable, str(MAIN_PY), "--step", "train_pa", "--PA_backbone", "bogus"] ) assert result.returncode != 0 + + +def test_training_defaults_match_opendpdv2_recipe(monkeypatch): + monkeypatch.setattr(sys, "argv", ["main.py"]) + + args = get_arguments() + + assert args.batch_size == 64 + assert args.batch_size_eval == 64 + assert args.n_epochs == 300 + assert args.frame_length == 200 + assert args.frame_stride == 1 + assert args.opt_type == "adamw" + assert args.loss_type == "l2" + assert args.seed == 0 + assert args.lr_schedule == 1 + assert args.lr == 5e-3 + assert args.lr_end == 5e-5 + assert args.decay_factor == 0.5 + assert args.patience == 5 diff --git a/tests/test_cuda_graph_frozen_dgru.py b/tests/test_cuda_graph_frozen_dgru.py new file mode 100644 index 0000000..18ee09a --- /dev/null +++ b/tests/test_cuda_graph_frozen_dgru.py @@ -0,0 +1,167 @@ +"""Correctness and fallback tests for frozen-DGRU CUDA-graph replay.""" + +import copy + +import pytest +import torch + +from backbones.cuda_graph_frozen_dgru import ( + clear_cuda_graph_frozen_dgru_cache, +) +from backbones.dgru import DGRU + + +requires_nvidia_cuda = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.version.cuda is None + or getattr(torch.version, "hip", None) is not None, + reason="requires NVIDIA CUDA", +) + + +def _build_frozen(device="cpu"): + model = DGRU(hidden_size=23, output_size=2, num_layers=1) + model = model.to(device).train() + model.reset_parameters() + for parameter in model.parameters(): + parameter.requires_grad_(False) + return model + + +def _is_graph_replay(output): + return "FrozenDGRUReplayFunction" in type(output.grad_fn).__name__ + + +def test_cpu_fallback_preserves_module_and_state_dict(monkeypatch): + clear_cuda_graph_frozen_dgru_cache() + monkeypatch.delenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", raising=False) + model = _build_frozen() + state_keys = tuple(model.state_dict()) + module_names = tuple(name for name, _ in model.named_modules()) + features = torch.randn(2, 7, 2, requires_grad=True) + initial_state = torch.zeros(1, 2, 23) + + output = model(features, initial_state) + output.square().mean().backward() + + assert not _is_graph_replay(output) + assert tuple(model.state_dict()) == state_keys + assert tuple(name for name, _ in model.named_modules()) == module_names + assert torch.isfinite(features.grad).all() + + +def test_trainable_dgru_uses_eager_fallback(): + model = DGRU(hidden_size=5, output_size=2, num_layers=1).train() + features = torch.randn(2, 4, 2, requires_grad=True) + output = model(features, torch.zeros(1, 2, 5)) + assert not _is_graph_replay(output) + output.sum().backward() + assert all(parameter.grad is not None for parameter in model.parameters()) + + +@requires_nvidia_cuda +def test_cuda_replay_is_bitwise_exact_and_state_dict_stable(monkeypatch): + torch.manual_seed(123) + clear_cuda_graph_frozen_dgru_cache() + model = _build_frozen("cuda") + state_before = copy.deepcopy(model.state_dict()) + module_names = tuple(name for name, _ in model.named_modules()) + initial_state = torch.randn(1, 4, 23, device="cuda") * 0.1 + + for _ in range(2): + features = torch.randn(4, 37, 2, device="cuda") + output_gradient = torch.randn(4, 37, 2, device="cuda") + monkeypatch.setenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", "1") + reference_input = features.clone().requires_grad_() + reference_output = model(reference_input, initial_state) + reference_output.backward(output_gradient) + + monkeypatch.delenv( + "OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", raising=False + ) + replay_input = features.clone().requires_grad_() + replay_output = model(replay_input, initial_state) + replay_output.backward(output_gradient) + torch.cuda.synchronize() + + assert _is_graph_replay(replay_output) + assert torch.equal(replay_output, reference_output) + assert torch.equal(replay_input.grad, reference_input.grad) + + assert tuple(model.state_dict()) == tuple(state_before) + assert tuple(name for name, _ in model.named_modules()) == module_names + for name, value in model.state_dict().items(): + assert torch.equal(value, state_before[name]) + + +@requires_nvidia_cuda +def test_busy_replay_and_opt_out_fall_back_safely(monkeypatch): + torch.manual_seed(321) + clear_cuda_graph_frozen_dgru_cache() + monkeypatch.delenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", raising=False) + model = _build_frozen("cuda") + initial_state = torch.zeros(1, 2, 23, device="cuda") + + first_input = torch.randn(2, 17, 2, device="cuda", requires_grad=True) + first_output = model(first_input, initial_state) + assert _is_graph_replay(first_output) + + # The first graph still owns its training reserve until backward. A + # second in-flight call must use independent eager buffers. + second_input = torch.randn(2, 17, 2, device="cuda", requires_grad=True) + second_output = model(second_input, initial_state) + assert not _is_graph_replay(second_output) + second_output.sum().backward() + first_output.sum().backward() + + monkeypatch.setenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", "true") + opt_out_input = torch.randn(2, 17, 2, device="cuda", requires_grad=True) + opt_out_output = model(opt_out_input, initial_state) + assert not _is_graph_replay(opt_out_output) + opt_out_output.sum().backward() + assert torch.isfinite(first_input.grad).all() + assert torch.isfinite(second_input.grad).all() + assert torch.isfinite(opt_out_input.grad).all() + + +@requires_nvidia_cuda +def test_autocast_and_eval_use_eager_fallback(monkeypatch): + clear_cuda_graph_frozen_dgru_cache() + monkeypatch.delenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", raising=False) + model = _build_frozen("cuda") + initial_state = torch.zeros(1, 2, 23, device="cuda") + features = torch.randn(2, 7, 2, device="cuda", requires_grad=True) + + with torch.autocast("cuda", dtype=torch.float16): + autocast_output = model(features, initial_state) + assert not _is_graph_replay(autocast_output) + + noncontiguous = torch.randn( + 2, 2, 7, device="cuda", requires_grad=True + ).transpose(1, 2) + noncontiguous_output = model(noncontiguous, initial_state) + assert not _is_graph_replay(noncontiguous_output) + + with torch.no_grad(): + no_grad_output = model(features, initial_state) + assert not no_grad_output.requires_grad + + model.eval() + eval_output = model(features, initial_state) + assert not _is_graph_replay(eval_output) + + +@requires_nvidia_cuda +def test_higher_order_gradient_fails_loudly(monkeypatch): + clear_cuda_graph_frozen_dgru_cache() + monkeypatch.delenv("OPENDPD_DISABLE_CUDA_GRAPH_FROZEN_DGRU", raising=False) + model = _build_frozen("cuda") + features = torch.randn(1, 3, 2, device="cuda", requires_grad=True) + initial_state = torch.zeros(1, 1, 23, device="cuda") + + with pytest.raises(RuntimeError, match="first-order training only"): + torch.autograd.grad( + model(features, initial_state).sum(), + features, + create_graph=True, + ) diff --git a/tests/test_cuda_graph_training.py b/tests/test_cuda_graph_training.py new file mode 100644 index 0000000..594e0c2 --- /dev/null +++ b/tests/test_cuda_graph_training.py @@ -0,0 +1,236 @@ +"""Tests for guarded whole-cascade CUDA-graph training.""" + +import contextlib +import copy +import io + +import pytest +import torch +from torch.utils.data import DataLoader, TensorDataset + +import modules.cuda_graph_training as graph_training +from backbones.cuda_graph_frozen_dgru import disable_cuda_graph_frozen_dgru +from models import CascadedModel, CoreModel +from modules.cuda_graph_training import ( + clear_cuda_graph_training_cache, + force_clean_eager_tres, + try_cuda_graph_training_step, +) +from modules.train_funcs import net_train + + +requires_nvidia_cuda = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.version.cuda is None + or getattr(torch.version, "hip", None) is not None, + reason="requires NVIDIA CUDA", +) + + +def _build_cascade(device="cpu"): + with contextlib.redirect_stdout(io.StringIO()): + dpd = CoreModel( + 2, 15, 1, "tres_deltagru", thx=0.01, thh=0.05 + ) + pa = CoreModel(2, 23, 1, "dgru") + cascade = CascadedModel(dpd, pa) + cascade.freeze_pa_model() + return cascade.to(device).train() + + +def _trainable(module): + return tuple( + parameter for parameter in module.parameters() + if parameter.requires_grad + ) + + +def test_cpu_opt_in_falls_back_cleanly_and_restores_dispatch(): + torch.manual_seed(301) + clear_cuda_graph_training_cache() + reference = _build_cascade() + candidate = copy.deepcopy(reference) + state_keys = tuple(candidate.state_dict()) + module_names = tuple(name for name, _ in candidate.named_modules()) + features = torch.randn(4, 5, 2) + targets = torch.randn(4, 5, 2) + loader = DataLoader( + TensorDataset(features, targets), batch_size=2, shuffle=False + ) + reference_optimizer = torch.optim.AdamW(_trainable(reference), lr=5e-4) + candidate_optimizer = torch.optim.AdamW(_trainable(candidate), lr=5e-4) + criterion = torch.nn.MSELoss() + + reference_log = {} + with force_clean_eager_tres(reference): + net_train( + reference_log, reference, loader, reference_optimizer, + criterion, 200.0, torch.device("cpu") + ) + candidate_log = {} + original_dispatch = candidate.dpd_model.backbone.rnn.use_triton + observed_dispatch = [] + hook = candidate.dpd_model.backbone.rnn.register_forward_pre_hook( + lambda recurrent, _inputs: observed_dispatch.append( + recurrent.use_triton + ) + ) + net_train( + candidate_log, candidate, loader, candidate_optimizer, + criterion, 200.0, torch.device("cpu"), + cuda_graph_training=True, + ) + hook.remove() + + assert candidate_log == reference_log + assert observed_dispatch and not any(observed_dispatch) + assert candidate.dpd_model.backbone.rnn.use_triton is original_dispatch + assert tuple(candidate.state_dict()) == state_keys + assert tuple(name for name, _ in candidate.named_modules()) == module_names + for expected, actual in zip( + reference.parameters(), candidate.parameters(), strict=True + ): + assert torch.equal(actual, expected) + + +def _snapshot(module, parameters, optimizer, output, loss): + return { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "gradients": [ + parameter.grad.detach().clone() for parameter in parameters + ], + "parameters": [parameter.detach().clone() for parameter in parameters], + "optimizer": [ + { + key: ( + value.detach().clone() + if torch.is_tensor(value) else value + ) + for key, value in optimizer.state[parameter].items() + } + for parameter in parameters + ], + "state_keys": tuple(module.state_dict()), + } + + +def _run_eager_trajectory(module, batches, targets): + parameters = _trainable(module) + optimizer = torch.optim.AdamW(parameters, lr=5e-4) + criterion = torch.nn.MSELoss() + output = loss = None + for features, target in zip(batches, targets, strict=True): + optimizer.zero_grad(set_to_none=True) + with force_clean_eager_tres(module), disable_cuda_graph_frozen_dgru(): + output = module(features) + loss = criterion(output, target) + loss.backward() + torch.nn.utils.clip_grad_norm_(parameters, 200.0) + optimizer.step() + return _snapshot(module, parameters, optimizer, output, loss) + + +def _run_graph_trajectory(module, batches, targets): + parameters = _trainable(module) + optimizer = torch.optim.AdamW(parameters, lr=5e-4) + criterion = torch.nn.MSELoss() + output = loss = None + original_dispatch = module.dpd_model.backbone.rnn.use_triton + for features, target in zip(batches, targets, strict=True): + result = try_cuda_graph_training_step( + module, features, target, criterion, parameters + ) + assert result is not None + assert module.dpd_model.backbone.rnn.use_triton is original_dispatch + output = result.output.clone() + loss = result.loss + torch.nn.utils.clip_grad_norm_(parameters, 200.0) + optimizer.step() + return _snapshot(module, parameters, optimizer, output, loss) + + +def _max_abs(left, right): + return max( + (first - second).abs().max().item() + for first, second in zip(left, right, strict=True) + ) + + +def _optimizer_max_abs(left, right): + differences = [] + for first, second in zip(left, right, strict=True): + assert first.keys() == second.keys() + for key in first: + if torch.is_tensor(first[key]): + differences.append( + (first[key] - second[key]).abs().max().item() + ) + else: + assert first[key] == second[key] + return max(differences, default=0.0) + + +@requires_nvidia_cuda +def test_cuda_ten_step_drift_is_within_eager_repeat_envelope(monkeypatch): + torch.manual_seed(401) + clear_cuda_graph_training_cache() + base = _build_cascade("cuda") + eager_first = copy.deepcopy(base) + eager_second = copy.deepcopy(base) + graphed = copy.deepcopy(base) + for module in (eager_first, eager_second, graphed): + module.pa_model.backbone.rnn.flatten_parameters() + generator = torch.Generator(device="cuda").manual_seed(402) + batch_sizes = [4, 3] * 5 + batches = [ + torch.randn(size, 17, 2, device="cuda", generator=generator) + for size in batch_sizes + ] + targets = [ + torch.randn(size, 17, 2, device="cuda", generator=generator) + for size in batch_sizes + ] + + first = _run_eager_trajectory(eager_first, batches, targets) + second = _run_eager_trajectory(eager_second, batches, targets) + capture_count = 0 + original_init = graph_training._CapturedTrainingStep.__init__ + + def counted_init(instance, *args, **kwargs): + nonlocal capture_count + capture_count += 1 + original_init(instance, *args, **kwargs) + + monkeypatch.setattr( + graph_training._CapturedTrainingStep, "__init__", counted_init + ) + candidate = _run_graph_trajectory(graphed, batches, targets) + torch.cuda.synchronize() + + assert candidate["state_keys"] == first["state_keys"] + assert capture_count == 2 + comparisons = ( + ( + (first["output"] - candidate["output"]).abs().max().item(), + (first["output"] - second["output"]).abs().max().item(), + ), + ( + abs((first["loss"] - candidate["loss"]).item()), + abs((first["loss"] - second["loss"]).item()), + ), + ( + _max_abs(first["gradients"], candidate["gradients"]), + _max_abs(first["gradients"], second["gradients"]), + ), + ( + _max_abs(first["parameters"], candidate["parameters"]), + _max_abs(first["parameters"], second["parameters"]), + ), + ( + _optimizer_max_abs(first["optimizer"], candidate["optimizer"]), + _optimizer_max_abs(first["optimizer"], second["optimizer"]), + ), + ) + for candidate_error, eager_repeat_error in comparisons: + assert candidate_error <= eager_repeat_error diff --git a/tests/test_metric_configuration.py b/tests/test_metric_configuration.py new file mode 100644 index 0000000..faea388 --- /dev/null +++ b/tests/test_metric_configuration.py @@ -0,0 +1,37 @@ +from types import SimpleNamespace + +import numpy as np + +from modules import train_funcs + + +def test_calculate_metrics_uses_dataset_sample_rate(monkeypatch): + prediction = np.zeros((1, 8, 2), dtype=np.float32) + ground_truth = np.zeros_like(prediction) + args = SimpleNamespace( + input_signal_fs=983.04e6, + bw_main_ch=200e6, + n_sub_ch=5, + nperseg=8, + ) + evm_kwargs = {} + + monkeypatch.setattr(train_funcs.metrics, "NMSE", lambda *_: -1.0) + + def fake_evm(*_, **kwargs): + evm_kwargs.update(kwargs) + return -2.0 + + monkeypatch.setattr(train_funcs.metrics, "EVM", fake_evm) + monkeypatch.setattr(train_funcs.metrics, "ACLR", lambda *_, **__: (-3.0, -5.0)) + + result = train_funcs.calculate_metrics(args, {}, prediction, ground_truth) + + assert evm_kwargs["sample_rate"] == args.input_signal_fs + assert result == { + "NMSE": -1.0, + "EVM": -2.0, + "ACLR_L": -3.0, + "ACLR_R": -5.0, + "ACLR_AVG": -4.0, + } diff --git a/tests/test_training_pipeline_optimizations.py b/tests/test_training_pipeline_optimizations.py new file mode 100644 index 0000000..c48dc83 --- /dev/null +++ b/tests/test_training_pipeline_optimizations.py @@ -0,0 +1,221 @@ +"""Equivalence checks for the optimized data and epoch plumbing.""" + +from copy import deepcopy +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset + +from modules.data_collector import IQFrameDataset +from modules.train_funcs import net_eval, net_train +from project import Project + + +def _reference_frames(sequence, frame_length, stride): + num_frames = (len(sequence) - frame_length) // stride + 1 + return np.stack([ + sequence[index * stride:index * stride + frame_length] + for index in range(num_frames) + ]) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int32]) +@pytest.mark.parametrize("stride", [1, 2, 4]) +def test_vectorized_frames_match_reference_and_do_not_alias(dtype, stride): + sequence = np.arange(34, dtype=dtype).reshape(17, 2) + expected = _reference_frames(sequence, frame_length=5, stride=stride) + + actual = IQFrameDataset.get_frames( + sequence, frame_length=5, stride_length=stride + ) + + assert np.array_equal(actual, expected) + assert actual.flags.c_contiguous + assert not np.shares_memory(actual, sequence) + + sequence[...] = 0 + assert np.array_equal(actual, expected) + + +def test_frame_dataset_preserves_values_order_and_independent_storage(): + features = np.arange(40, dtype=np.float64).reshape(20, 2) / 7 + targets = -features.copy() + expected_features = torch.Tensor(_reference_frames(features, 6, 3)) + expected_targets = torch.Tensor(_reference_frames(targets, 6, 3)) + + dataset = IQFrameDataset(features, targets, frame_length=6, stride=3) + features[...] = 99 + targets[...] = 99 + + assert dataset.features.dtype == torch.float32 + assert dataset.targets.dtype == torch.float32 + assert torch.equal(dataset.features, expected_features) + assert torch.equal(dataset.targets, expected_targets) + + +class _FrozenTail(nn.Module): + """Tiny analogue of a trainable DPD followed by a frozen PA.""" + + def __init__(self): + super().__init__() + self.trainable = nn.Linear(3, 4) + self.frozen = nn.Linear(4, 2) + for parameter in self.frozen.parameters(): + parameter.requires_grad_(False) + + def forward(self, inputs): + return self.frozen(torch.tanh(self.trainable(inputs))) + + +def _reference_train( + log, net, dataloader, optimizer, criterion, grad_clip_val +): + """The pre-optimization epoch loop, retained only as a test oracle.""" + net.train() + losses = [] + for features, targets in dataloader: + features = features.to(torch.device("cpu")) + targets = targets.to(torch.device("cpu")) + optimizer.zero_grad() + loss = criterion(net(features), targets) + loss.backward() + if grad_clip_val != 0: + nn.utils.clip_grad_norm_(net.parameters(), grad_clip_val) + optimizer.step() + loss.detach() + losses.append(loss.item()) + log["loss"] = np.mean(losses) + return net + + +def _build_project_optimizer(net): + project = Project.__new__(Project) + project.opt_type = "adamw" + project.lr = 2e-3 + project.decay_factor = 0.5 + project.patience = 2 + project.lr_end = 1e-7 + return project.build_optimizer(net)[0] + + +@pytest.mark.parametrize("num_batches", [1, 4]) +def test_train_epoch_matches_original_updates_and_logged_loss(num_batches): + torch.manual_seed(71) + features = torch.randn(num_batches * 3, 3) + targets = torch.randn(num_batches * 3, 2) + dataloader = DataLoader( + TensorDataset(features, targets), batch_size=3, shuffle=False + ) + + reference_net = _FrozenTail() + optimized_net = deepcopy(reference_net) + reference_optimizer = torch.optim.AdamW( + reference_net.parameters(), lr=2e-3 + ) + optimized_optimizer = _build_project_optimizer(optimized_net) + + optimized_parameters = { + id(parameter) + for group in optimized_optimizer.param_groups + for parameter in group["params"] + } + assert optimized_parameters == { + id(parameter) for parameter in optimized_net.parameters() + if parameter.requires_grad + } + + reference_log = {} + optimized_log = {} + criterion = nn.MSELoss() + _reference_train( + reference_log, reference_net, dataloader, reference_optimizer, + criterion, grad_clip_val=0.7 + ) + net_train( + optimized_log, optimized_net, dataloader, optimized_optimizer, + criterion, grad_clip_val=0.7, device=torch.device("cpu") + ) + + assert optimized_log["loss"] == reference_log["loss"] + for reference, optimized in zip( + reference_net.parameters(), optimized_net.parameters(), strict=True + ): + assert torch.equal(optimized, reference) + + +def _reference_eval(log, net, dataloader, criterion): + net.eval() + losses = [] + prediction = [] + ground_truth = [] + with torch.no_grad(): + for features, targets in dataloader: + outputs = net(features) + loss = criterion(outputs, targets) + prediction.append(outputs.cpu()) + ground_truth.append(targets.cpu()) + losses.append(loss.item()) + log["loss"] = np.mean(losses) + return ( + torch.cat(prediction, dim=0).numpy(), + torch.cat(ground_truth, dim=0).numpy(), + ) + + +def test_inference_eval_matches_original_outputs_targets_and_loss(): + torch.manual_seed(101) + net = _FrozenTail() + features = torch.randn(11, 3) + targets = torch.randn(11, 2) + dataloader = DataLoader( + TensorDataset(features, targets), batch_size=4, shuffle=False + ) + criterion = nn.L1Loss() + + reference_log = {} + reference_prediction, reference_targets = _reference_eval( + reference_log, net, dataloader, criterion + ) + optimized_log = {} + _, optimized_prediction, optimized_targets = net_eval( + optimized_log, net, dataloader, criterion, torch.device("cpu") + ) + + assert optimized_log["loss"] == reference_log["loss"] + assert np.array_equal(optimized_prediction, reference_prediction) + assert np.array_equal(optimized_targets, reference_targets) + + +@pytest.mark.parametrize( + "device_type, expected", [("cpu", False), ("cuda", True)] +) +def test_cuda_loaders_pin_memory(monkeypatch, device_type, expected): + features = np.arange(48, dtype=np.float64).reshape(24, 2) / 10 + targets = features * 2 + + def fake_load_dataset(**_kwargs): + return ( + features, targets, + features[:12], targets[:12], + features[:12], targets[:12], + ) + + monkeypatch.setattr( + "modules.data_collector.load_dataset", fake_load_dataset + ) + project = Project.__new__(Project) + project.dataset_name = "unused" + project.step = "train_pa" + project.frame_length = 4 + project.frame_stride = 2 + project.batch_size = 3 + project.batch_size_eval = 2 + project.args = SimpleNamespace(nperseg=6) + project.device = torch.device(device_type) + + loaders, _ = project.build_dataloaders() + + assert all(loader.pin_memory is expected for loader in loaders) diff --git a/tests/test_tres_deltagru_optimized.py b/tests/test_tres_deltagru_optimized.py new file mode 100644 index 0000000..abdd20e --- /dev/null +++ b/tests/test_tres_deltagru_optimized.py @@ -0,0 +1,263 @@ +"""Correctness tests for the fused TRes-DeltaGRU recurrence.""" + +import contextlib +import copy +import io +from types import SimpleNamespace + +import pytest +import torch + +from models import CoreModel +from backbones.tres_deltagru import DeltaGRULayer +from backbones.triton_deltagru import triton +from modules.paths import gen_log_stat + + +requires_fused_cuda = pytest.mark.skipif( + not torch.cuda.is_available() or triton is None, + reason="requires CUDA and Triton", +) + + +def build_model(device, *, fused, thx=0.01, thh=0.05): + with contextlib.redirect_stdout(io.StringIO()): + model = CoreModel( + input_size=2, + hidden_size=15, + num_layers=1, + backbone_type="tres_deltagru", + thx=thx, + thh=thh, + ).to(device) + model.backbone.rnn.use_triton = fused + model.backbone.set_debug(0) + return model + + +def test_debug_setting_reaches_recurrent_layer(): + model = build_model("cpu", fused=False) + model.backbone.set_debug(1) + assert model.backbone.debug == 1 + assert model.backbone.rnn.debug == 1 + model.backbone.set_debug(0) + assert model.backbone.debug == 0 + assert model.backbone.rnn.debug == 0 + + +def test_h15_matches_the_approximately_1000_parameter_dpd_budget(): + model = build_model("cpu", fused=False, thx=0.0, thh=0.0) + + assert sum(parameter.numel() for parameter in model.parameters()) == 999 + + +def test_pa_delta_thresholds_are_recorded(): + model = build_model("cpu", fused=False, thx=0.0, thh=0.0) + args = SimpleNamespace( + step="train_pa", + n_epochs=2, + batch_size=2, + frame_length=3, + PA_backbone="tres_deltagru", + PA_hidden_size=15, + ) + + log = gen_log_stat(args, 0.0, model, None, 0) + + assert log["THX"] == 0.0 + assert log["THH"] == 0.0 + + +def test_opt_in_statistics_are_logged_and_reset(): + model = build_model("cpu", fused=False) + model.backbone.set_debug(1) + model(torch.randn(2, 3, 2)) + wrapper = SimpleNamespace( + dpd_model=model, + named_parameters=model.named_parameters, + ) + args = SimpleNamespace( + step="train_dpd", + n_epochs=2, + batch_size=2, + frame_length=3, + DPD_backbone="tres_deltagru", + DPD_hidden_size=15, + ) + + log = gen_log_stat(args, 0.0, wrapper, None, 0) + + assert log["THX"] == 0.01 + assert log["THH"] == 0.05 + assert {"SP_T_DX", "SP_T_DH", "SP_T_DV", "HW_PARAM"} <= log.keys() + assert model.backbone.rnn.statistics == { + "num_dx_zeros": 0, + "num_dx_numel": 0, + "num_dh_zeros": 0, + "num_dh_numel": 0, + } + + +@requires_fused_cuda +@pytest.mark.parametrize("thresholds", [(0.0, 0.0), (0.01, 0.05)]) +@pytest.mark.parametrize("sequence_length", [1, 17, 200]) +def test_fused_output_and_gradients_match_eager(thresholds, sequence_length): + torch.manual_seed(7) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + thx, thh = thresholds + reference = build_model("cuda", fused=False, thx=thx, thh=thh) + fused = build_model("cuda", fused=True, thx=thx, thh=thh) + fused.load_state_dict(copy.deepcopy(reference.state_dict())) + + reference_input = torch.randn( + 2, sequence_length, 2, device="cuda", requires_grad=True + ) + fused_input = reference_input.detach().clone().requires_grad_() + target = torch.randn(2, sequence_length, 2, device="cuda") + + reference_output = reference(reference_input) + reference_loss = torch.nn.functional.mse_loss(reference_output, target) + reference_loss.backward() + fused_output = fused(fused_input) + fused_loss = torch.nn.functional.mse_loss(fused_output, target) + fused_loss.backward() + + torch.testing.assert_close(fused_output, reference_output, rtol=3e-5, atol=2e-5) + torch.testing.assert_close(fused_input.grad, reference_input.grad, rtol=5e-4, atol=2e-6) + for (name, reference_parameter), (fused_name, fused_parameter) in zip( + reference.named_parameters(), fused.named_parameters(), strict=True + ): + assert name == fused_name + torch.testing.assert_close( + fused_parameter.grad, + reference_parameter.grad, + rtol=5e-4, + atol=2e-6, + msg=lambda message, parameter=name: f"{parameter}: {message}", + ) + + +@requires_fused_cuda +def test_fused_statistics_match_eager(): + torch.manual_seed(11) + reference = build_model("cuda", fused=False) + fused = build_model("cuda", fused=True) + fused.load_state_dict(copy.deepcopy(reference.state_dict())) + reference.backbone.set_debug(1) + fused.backbone.set_debug(1) + features = torch.randn(4, 17, 2, device="cuda") + + reference(features) + fused(features) + + for key in ("num_dx_zeros", "num_dx_numel", "num_dh_zeros", "num_dh_numel"): + reference_value = torch.as_tensor(reference.backbone.rnn.statistics[key]) + fused_value = torch.as_tensor(fused.backbone.rnn.statistics[key]) + torch.testing.assert_close(fused_value.cpu(), reference_value.cpu()) + + +@requires_fused_cuda +def test_custom_initial_state_gradients_match_eager(): + torch.manual_seed(13) + reference = DeltaGRULayer(6, 15, 1, thx=0.01, thh=0.05).cuda() + fused = copy.deepcopy(reference) + reference.use_triton = False + fused.use_triton = True + reference.debug = 0 + fused.debug = 0 + + reference_input = torch.randn(2, 17, 6, device="cuda", requires_grad=True) + fused_input = reference_input.detach().clone().requires_grad_() + + def make_state(width): + return (torch.randn(1, 2, width, device="cuda") * 0.1).requires_grad_() + + reference_states = [ + make_state(15), # x_p_0 uses x_p_length=max(input_size, hidden_size) + make_state(15), + make_state(15), + make_state(15), + make_state(45), + ] + fused_states = [ + state.detach().clone().requires_grad_() for state in reference_states + ] + output_gradient = torch.randn(2, 17, 15, device="cuda") + + reference_output = reference(reference_input, *reference_states) + (reference_output * output_gradient).sum().backward() + fused_output = fused(fused_input, *fused_states) + (fused_output * output_gradient).sum().backward() + + torch.testing.assert_close(fused_output, reference_output, rtol=3e-5, atol=2e-5) + torch.testing.assert_close(fused_input.grad, reference_input.grad, rtol=5e-4, atol=2e-6) + for reference_state, fused_state in zip( + reference_states, fused_states, strict=True + ): + torch.testing.assert_close( + fused_state.grad, reference_state.grad, rtol=5e-4, atol=2e-6 + ) + + +@requires_fused_cuda +def test_inference_does_not_save_training_workspace(): + model = build_model("cuda", fused=True).eval() + features = torch.randn(4, 512, 2, device="cuda") + with torch.inference_mode(): + output = model(features) + assert output.shape == (4, 512, 2) + assert not output.requires_grad + assert torch.isfinite(output).all() + + +@requires_fused_cuda +def test_cuda_autocast_uses_equivalent_eager_fallback(): + torch.manual_seed(17) + reference = build_model("cuda", fused=False) + candidate = build_model("cuda", fused=True) + candidate.load_state_dict(copy.deepcopy(reference.state_dict())) + features = torch.randn(2, 17, 2, device="cuda") + + with torch.autocast("cuda", dtype=torch.float16): + reference_output = reference(features) + candidate_output = candidate(features) + probe = torch.empty(1, 1, 6, device="cuda") + assert not candidate.backbone.rnn._can_use_triton(probe) + + torch.testing.assert_close(candidate_output, reference_output, rtol=0, atol=0) + + +@requires_fused_cuda +def test_quantization_aware_model_uses_eager_fallback(): + from quant.quant_envs import AttrDict, Base_GRUQuantEnv + + float_model = build_model("cpu", fused=True) + with contextlib.redirect_stdout(io.StringIO()): + environment = Base_GRUQuantEnv( + float_model, + AttrDict( + n_bits_w=16, + n_bits_a=16, + pretrained_model="", + quant_dir_label="", + ), + ) + quantized_model = environment.q_model.cuda() + probe = torch.empty(1, 1, 6, device="cuda") + assert not quantized_model.backbone.rnn._can_use_triton(probe) + + features = torch.randn(2, 4, 2, device="cuda") + output = quantized_model(features) + output.square().mean().backward() + assert torch.isfinite(output).all() + + +@requires_fused_cuda +def test_higher_order_gradient_fails_loudly(): + model = build_model("cuda", fused=True) + features = torch.randn(1, 2, 2, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="first-order training only"): + torch.autograd.grad( + model(features).sum(), features, create_graph=True + )