StatBaselineTrainer

class ts2.models.single_session.StatBaselineTrainer(cfg, rank, world_size)

Bases: ts2.ts2_eval_trainer.TS2EvalTrainer

TS2 trainer for the zero-parameter statistical baselines.

The model is fit in link_datasets, so there is no optimizer, train_epoch is a no-op and val reports the fit’s own selection score.

setup_optimizers(ckpt)

Configure the AdamW optimizer and OneCycleLR scheduler.

Weight decay skips 1-D params and the names in cfg.no_weight_decay. Override this to use a different optimizer or scheduler.

Parameters:

ckpt (dict | None) – Checkpoint dictionary. Optimizer/scheduler resumption is not yet supported.

Returns:

(optimizer, scheduler)

Return type:

tuple

train_epoch()

Run one training epoch.

Iterates over the train loader, computes predictions and loss, and updates model parameters via optimizer and scheduler. Logs loss per step if log_train_step is enabled, otherwise logs the epoch average.

link_model(model, ckpt)

Link datasets to the model.

Attaches the configured transforms and the model input_fn to the train/val/test dataset transform pipelines, then calls model.link_datasets to register the datasets with the model.

Parameters:

• model (BaseModel) – The instantiated model to link.

• ckpt (dict | None) – Checkpoint dictionary, passed through for subclass use.

val_epoch()

Report the fit’s own val score instead of walking the val loader.

The fit already scored its selection on the val split with the same metric, so the walk would only re-measure it at a denser stride. No bps, and nothing to early-stop, so best_model stays unset and test scores the fitted model directly.