sage.factory.training

Sage training loop.

SageVanillaTraining

The standard Sage training loop. Compiles nothing — the neural network model should be compiled externally with torch.compile(model, ...) before being passed in, which is where the meaningful speedup comes from.

Auto-multibanding: if the signal sampler has a selector attribute (set when multiband_mode='worst_case' is used), the noise batch is automatically projected to the coarse grid before injection — the user does not need a separate noise wrapper.

GWBatch tracking: after combining signal + noise, the joint tensor is wrapped in a GWBatch that carries the current grid type and whitening state through the preprocessing pipeline. Pipeline violations (e.g. IFFT on a non-uniform grid, multirate sampling before TD conversion) raise PipelineError immediately when the offending module is called.

The model receives batch.to_network_input() — a real float32 tensor whose shape is (B, D, T) for TD pipelines and (B, 2·D, N_coarse) for FD_COARSE pipelines.

Classes

SageVanillaTraining

Standard Sage training loop.

Module Contents

class SageVanillaTraining(signal_sampler, noise_sampler, processor, model, loss_function, optimiser, scheduler, scaler, num_iterations, num_epochs, scheduler_mode='batch', callbacks=None, aux_losses=None, balancer=None)[source]

Bases: torch.nn.Module

Standard Sage training loop.

Each epoch iterates over num_iterations batches:

Draw a signal batch (S, D, F/T) and a noise batch (B, D, F/T) from the respective samplers.
If the signal sampler uses worst-case multibanding, apply the selector to the noise batch automatically (no user action required).
Scatter-inject signal waveforms at random positions in the noise batch.
Wrap the combined tensor in a GWBatch carrying the current grid state.
Pass through the preprocessing pipeline (whitening, optional IFFT and multirate sampling).
Forward pass through the compiled model, loss, backpropagation.

Parameters:

signal_sampler (nn.Module) – Waveform signal sampler (e.g. IMRPhenomXAS_NRTidalv3). If it exposes output_state and selector attributes (set automatically when multiband_mode='worst_case' is used), the training loop configures itself accordingly.
noise_sampler (nn.Module) – Noise sampler (e.g. MemmapNoiseSampler). No changes needed for multibanding — the selector is applied internally.
processor (nn.Module) – Preprocessing pipeline (e.g. Preprocessor wrapping FiducialWhitening and optionally MultirateSampler).
model (nn.Module) – The neural network. Should be compiled externally with torch.compile(model) before being passed here.
loss_function (nn.Module) – Loss function returning a stacked component tensor.
optimiser (torch.optim.Optimizer)
scheduler (torch.optim.lr_scheduler._LRScheduler)
scaler (torch.amp.GradScaler or None) – AMP gradient scaler. Pass None to disable AMP.
num_iterations (int) – Gradient steps per epoch.
num_epochs (int) – Total epochs (pre-allocates the loss tracking tensor).
scheduler_mode (str) – "batch" to step the scheduler every batch, "epoch" once per epoch.
state (Initialize internal Module)
ScriptModule. (shared by both nn.Module and)