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