locpix.img_processing.training.train.train_loop#

train_loop(epochs, model, optimiser, train_loader, val_loader, loss_fn, device, model_path)[source]#

Function defining the training loop

Parameters:
  • epochs (int) – Number of epochs to train for

  • optimiser (torch.optim optimiser) – Optimiser that controls training of the model

  • model (torch geometric model) – Model that is going to be trained

  • train_loader (torch dataloader) – Dataloader for the training dataset

  • val_loader (torch dataloader) – Dataloader for the validation dataset

  • loss_fn (loss function) – Loss function calculate loss between train and output

  • device (gpu or cpu) – Device to train the model on

  • model_path (string) – Where to save model to