TS1EvalTrainer#
- class ts1.TS1EvalTrainer(cfg, rank, world_size)[source]#
Bases:
ts1.ts1_test_mixin.TS1TestMixin,core.trainer.BaseTrainerStandardized evaluation Trainer for Task Suite 1 (decoding).
This class defines the base Trainer for benchmark decoding tasks. This interface abstracts away the common boilerplate code for evaluating models on the benchmark, and allows for easy customization of the training loop. Several methods related to standardized evaluation are declared final and should not be overridden. Please refer to
POYOEvalTrainerfor 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.
predict: Given a batch of data, including inputs as defined by the model’s input_fn and information about the target (timestamps and mask), generate the model’s predictions.
See Task Suite 1: Behavior 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.
- 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.
- link_model(model, ckpt)[source]#
Link datasets to the model and configure the readout head.
Attaches model-specific transforms and input_fn to the train/val/test dataset pipelines, then calls
model.link_datasetsandmodel.configure_readoutto wire up the readout head for the current task.
- setup_finetuning()[source]#
Set up the finetuning strategy from config.
Instantiates and initializes the finetuning strategy if
finetuning.strategyis specified in the config. Setsself.ft_strategytoNoneotherwise.
- loss(pred, X, y)[source]#
Compute the task loss.
Override this to add auxiliary losses or custom loss logic.
- Parameters:
pred – Model predictions.
X – Input batch dict.
y – Target batch dict.
- Returns:
Scalar loss value.
- Return type:
- 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_stepis enabled, otherwise logs the epoch average.
- predict(X, target_timestamps=None, target_mask=None)[source]#
Generate model predictions for a batch.
Override this if the model requires target timestamps or mask at inference time (e.g. for masked prediction models).
- Parameters:
X – Input batch dict containing
model_inputs.target_timestamps – Optional target timestamps tensor.
target_mask – Optional boolean mask tensor for the targets.
- Returns:
Model output predictions.
- Return type:
- 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.