TS2EvalTrainer#

class ts2.TS2EvalTrainer(cfg, rank, world_size)[source]#

Bases: ts2.ts2_test_mixin.TS2TestMixin, core.trainer.BaseTrainer

Standardized evaluation Trainer for Task Suite 2 (neural prediction).

This class defines the base Trainer for benchmark neural prediction tasks. This interface abstracts away the common boilerplate code for evaluating models on the benchmark, and allows for easy customization of the training loop. Both tasks reconstruct held-out spike counts stripped from the input before input_fn, and only those entries are scored. Several methods related to standardized evaluation are declared final and should not be overridden. Please refer to LFADSEvalTrainer for a template for customizing the training logic.

Available to customize (override these methods in a custom Trainer, if needed):

  • setup: Set up the trainer. This includes setting up the tasks, data loaders, model, optimizer, and scheduler.

  • setup_model: Instantiate and initialize the model.

  • link_model: Link the datasets to the model.

  • setup_finetuning: Set up the finetuning strategy.

  • setup_optimizers: Configure the optimizer (and scheduler, if applicable).

  • loss: Define the loss function.

  • train_epoch: Define the training loop. This is flexible and allows for custom training logic, such as adding gradient accumulation, curriculum learning, etc.

  • training_step: Compute the loss for one batch. This is where a model that corrupts its own input applies its masker.

  • predict: Given a batch of data, including inputs as defined by the model’s input_fn and any masking information, generate the model’s predictions.

See Task Suite 2: Neural Activity Prediction for more details on the interface.

setup(ckpt)[source]#

Set up all trainer components.

Called automatically by the base constructor. Initializes tasks, data loaders, model, DDP wrapping, optimizer, and finetuning strategy in order.

Parameters:

ckpt (dict | None) – Checkpoint dictionary loaded from disk, or None for a fresh run.

setup_model(ckpt)[source]#

Instantiate and register the model.

Instantiates the model from the Hydra config and registers it for checkpointing. Override this to load pretrained weights or modify the model architecture.

Parameters:

ckpt (dict | None) – Checkpoint dictionary, available for loading pretrained weights.

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.

setup_finetuning()[source]#

Set up the finetuning strategy from config.

Instantiates and initializes the finetuning strategy if finetuning.strategy is specified in the config. Sets self.ft_strategy to None otherwise.

loss(pred, target, mask=None)[source]#

Compute the reconstruction loss.

Parameters:

  • pred – Model predictions.

  • target – Ground truth targets.

  • mask – Optional boolean mask. If provided, loss is averaged only over masked positions.

Returns:

Scalar loss value.

Return type:

torch.Tensor

train_epoch()[source]#

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.

predict(X, mask=None, mask_timestamps=None, mask_units=None)[source]#

Generate model predictions for a batch.

Override this if the model requires masking information at inference time (e.g. for masked autoencoder models).

Parameters:

  • X – Input batch dict containing model_inputs.

  • mask – Optional boolean mask for masked prediction.

  • mask_timestamps – Optional timestamps for masked positions.

  • mask_units – Optional unit indices for masked positions.

Returns:

Model output predictions.

Return type:

torch.Tensor

training_step(X, y)[source]#

Compute loss for one training batch.

Override this to customize masking, auxiliary losses, or any model-specific training logic. The base implementation is a plain forward pass against the benchmark-fixed target.

Parameters:

  • X (dict) – Input batch dict containing model_inputs.

  • y (dict) – Target batch dict containing values (benchmark-fixed spike counts).

Returns:

Scalar loss.

Return type:

torch.Tensor

setup_optimizers(ckpt)[source]#

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

log_training_results()[source]#

Print a Rich table summarizing the best metrics at the end of training.

val_epoch()[source]#

Perform a validation epoch.

Called every val.every_n_epochs epochs, and at the end of training.

Returns:

bool, True if early stopping should be triggered, False (or falsey) otherwise.

Return type:

early_stop