kosmos.ml.sl_trainer

Module Attributes

DEVICE = 'cpu'

Classes

class SLTrainer(model: kosmos.ml.models.model.Model, optimizer_config: kosmos.ml.config.factories.optimizer.OptimizerConfig, lr_scheduler_config: kosmos.ml.config.factories.lr_scheduler.LearningRateSchedulerConfig | None, loss_config: kosmos.ml.config.factories.loss.LossConfig, max_grad_norm: float | None)

Trainer for supervised learning classification tasks.

Initialize a supervised learning trainer.

Parameters:

  • model (Model) – The model to train/evaluate.

  • optimizer_config (OptimizerConfig) – Optimizer configuration.

  • lr_scheduler_config (LearningRateSchedulerConfig) – Learning rate scheduler configuration.

  • loss_config (LossConfig) – Loss configuration.

  • max_grad_norm (float | None) – Maximum gradient norm.


Methods

train(num_epochs: int, dataloader: torch.utils.data.DataLoader, fl_round: int | None = None, node: kosmos.topology.node.Node | None = None) → collections.abc.Iterator[kosmos.ml.sl_result.SLTrainIterationResult]

Train the model on the given train data.

Parameters:

  • num_epochs (int) – Number of epochs to run.

  • dataloader (DataLoader) – DataLoader providing the train data.

  • fl_round (int | None) – Federated learning round index to attach to the results. Defaults to None.

  • node (Node | None) – Node corresponding to this iteration to attach to the results. Defaults to None.

Returns:

An iterator yielding one train result per epoch.

Return type:

Iterator[SLTrainIterationResult]

test(dataloader: torch.utils.data.DataLoader) → kosmos.ml.sl_result.SLTestIterationResult

Evaluate the model on the given test dataloader.

Parameters:

dataloader (DataLoader) – DataLoader providing the test data.

Returns:

Result of the test iteration.

Return type:

SLTestIterationResult