neuralop.training.Trainer

class neuralop.training.Trainer(*, model, n_epochs, wandb_log=True, device=None, amp_autocast=False, data_processor=None, callbacks=None, log_test_interval=1, log_output=False, use_distributed=False, verbose=False)[source]

Methods

evaluate(loss_dict, data_loader[, log_prefix])

Evaluates the model on a dictionary of losses

train(train_loader, test_loaders, optimizer, ...)

Trains the given model on the given datasets. params: train_loader: torch.utils.data.DataLoader training dataloader test_loaders: dict[torch.utils.data.DataLoader] testing dataloaders optimizer: torch.optim.Optimizer optimizer to use during training optimizer: torch.optim.lr_scheduler learning rate scheduler to use during training training_loss: training.losses function cost function to minimize eval_losses: dict[Loss] dict of losses to use in self.eval().

train(train_loader, test_loaders, optimizer, scheduler, regularizer, training_loss=None, eval_losses=None)[source]

Trains the given model on the given datasets. params: train_loader: torch.utils.data.DataLoader

training dataloader

test_loaders: dict[torch.utils.data.DataLoader]

testing dataloaders

optimizer: torch.optim.Optimizer

optimizer to use during training

optimizer: torch.optim.lr_scheduler

learning rate scheduler to use during training

training_loss: training.losses function

cost function to minimize

eval_losses: dict[Loss]

dict of losses to use in self.eval()

evaluate(loss_dict, data_loader, log_prefix='')[source]

Evaluates the model on a dictionary of losses

Parameters:

loss_dictdict of functions

each function takes as input a tuple (prediction, ground_truth) and returns the corresponding loss

data_loaderdata_loader to evaluate on

log_prefixstr, default is ‘’

if not ‘’, used as prefix in output dictionary

Returns:

errorsdict

dict[f’{log_prefix}_{loss_name}] = loss for loss in loss_dict