Which one (train loss or validation loss) should I take consider to evalute a model at the time of k-fold cross validation?

Question

I want to be sure about train loss and validation loss to evaluate a model. My target is to train a model using (k-1) folds and then test the model with the rest of the 1 fold dataset and want to save the predictions.

Say, the MNIST digits dataset contains 70,000 (60,000 train and 10,000 test) images. I want to use k-fold cross-validation on the whole dataset (70,000 images).

Now, I trained a model using (k-1) folds based on the train loss as the monitor. I am saving (at checkpoints) top 3 models based on the performance (train loss). So, if I have 3 epochs then I will do have 9 (3*3) models. Later, the rest of one fold dataset was used to test the model prediction/inference. However, I know, usually, people do use validation loss to evaluate the best model.

K-fold code sample

k_folds=5
epochs=3

kfold=KFold(n_splits=k_folds,shuffle=True)
dataset = prepare_data()
model = LightningMNIST(lr_rate=0.01)

for fold, (train_idx, val_idx) in enumerate(kfold.split(dataset)):
  train_subsampler = torch.utils.data.SubsetRandomSampler(train_idx)
  val_subsampler = torch.utils.data.SubsetRandomSampler(val_idx)
  train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, sampler=train_subsampler)
  val_loader = torch.utils.data.DataLoader(dataset, batch_size=64, sampler=val_subsampler)

  model.apply(reset_weights) # resetting model weight for every fold

  
  early_stopping = EarlyStopping('train_loss', mode='min', patience=5)
  model_checkpoint = ModelCheckpoint(dirpath=model_path+'mnist_{epoch}-{train_loss:.2f}',
                                          monitor='train_loss', mode='min', save_top_k=3)
  trainer = pl.Trainer(max_epochs=epochs, profiler=False, callbacks = [model_checkpoint],default_root_dir=model_path) 
  trainer.fit(model, train_dataloader=train_loader)
  trainer.test(test_dataloaders=val_loader, ckpt_path=None)

Now, my concern is, am I doing the wrong k-fold (my logic is wrong)? Should I use the 1 fold dataset for validation and use the validation loss to check the model performance?

Note: I have used pytorch-lightning.

Answer 1

Typically k-fold is not used in neural networks because it's more expensive. When dataset is large enough, one validation set is more economic, and you end up with a single final model.

The test set (i.e. 10K samples) is not to be touched until your model is ready. So, you'll use your 60K samples for train+validation. A 15-20 % of the 60K samples can be chosen as validation for example. And, for the early stopping, you should look at validation loss, because that's more meaningful in terms of generalization performance. Once your model finishes the training, the performance is reported on the test set.