Trainer

Class to train a neural network using a Pytorch optimization algorithm.

The Trainer class is imported using the following command:

>>> from UQpy.scientific_machine_learning import Trainer

Methods

class Trainer(model, optimizer, loss_function=MSELoss(), scheduler=None)[source]

Prepare to train a neural network

Parameters:

  • model (Module) – Neural Network model to be trained

  • optimizer (Optimizer) – Optimization algorithm used to update model parameters

  • loss_function (Module) – Function used to compute loss during training

  • scheduler (Union[LRScheduler, list, None]) – Scheduler used to adjust the learning rate of the optimizer. Schedulers may be chained together by creating a list of schedulers

run(train_data=None, test_data=None, epochs=100, tolerance=0.0)[source]

Run the optimizer algorithm to learn the parameters of the model that fit train_data

Parameters:

  • train_data (Optional[DataLoader]) – Data used to compute model loss

  • test_data (Optional[DataLoader]) – Data used to validate the performance of the model

  • epochs (int) – Maximum number of epochs to run the optimizer for

  • tolerance (float) – Optimization terminates early if average training loss is below tolerance. Default: \(0.0\)

Raises:

RuntimeError – If neither train_data nor test_data is provided, a RuntimeError occurs.

Attributes

Trainer.history: dict

Record of the loss during training and validation. Note if training ends early there may be NaN values, as the histories are initialized with NaN.

  • history["train_loss"] contains training history as a torch.Tensor.

  • history["test_loss"] contains testing history as a torch.Tensor.

Examples