How to record each fold's validation loss during cross-validation in Optuna?

I am using Toshihiko Yanase's code for doing cross validation on my hyperparameter optimizer with Optuna. Here is the code that I am using:

def objective(trial, train_loader, valid_loader):

    # Remove the following line.
    # train_loader, valid_loader = get_mnist()

    ...

    return accuracy


def objective_cv(trial):

    # Get the MNIST dataset.
    dataset = datasets.MNIST(DIR, train=True, download=True, transform=transforms.ToTensor())

    fold = KFold(n_splits=3, shuffle=True, random_state=0)
    scores = []
    for fold_idx, (train_idx, valid_idx) in enumerate(fold.split(range(len(dataset)))):
        train_data = torch.utils.data.Subset(dataset, train_idx)
        valid_data = torch.utils.data.Subset(dataset, valid_idx)

        train_loader = torch.utils.data.DataLoader(
            train_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )
        valid_loader = torch.utils.data.DataLoader(
            valid_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )

        accuracy = objective(trial, train_loader, valid_loader)
        scores.append(accuracy)
    return np.mean(scores)


study = optuna.create_study(direction="maximize")
study.optimize(objective_cv, n_trials=20, timeout=600)

Unfortunately, using the code this way, it does not record each folds val loss to the Optuna dashboard. Is there a way to record each folds val loss to the Optuna dashboard?

Solution

Each splits validation loss can be recorded in the system_attrs of the Trial object of the current trial. The system_attrs can be seen in the dashboard under the respective trial as you wished.

The modified code having the desired functionality is:

def objective(trial, train_loader, valid_loader):

    # Remove the following line.
    # train_loader, valid_loader = get_mnist()

    ...

    return accuracy


def objective_cv(trial):

    # Get the MNIST dataset.
    dataset = datasets.MNIST(DIR, train=True, download=True, transform=transforms.ToTensor())

    fold = KFold(n_splits=3, shuffle=True, random_state=0)
    scores = []
    trial.set_system_attr("Val loss of fold",[])   #to record each individual final loss of the current fold
    for fold_idx, (train_idx, valid_idx) in enumerate(fold.split(range(len(dataset)))):
        train_data = torch.utils.data.Subset(dataset, train_idx)
        valid_data = torch.utils.data.Subset(dataset, valid_idx)

        train_loader = torch.utils.data.DataLoader(
            train_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )
        valid_loader = torch.utils.data.DataLoader(
            valid_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )

        accuracy = objective(trial, train_loader, valid_loader)
        scores.append(accuracy)
        trial.set_system_attr("Val loss of fold",trial.system_attrs["Val loss of fold"]+[accuracy]) #here is the objective value is added to the record
    return np.mean(scores)


study = optuna.create_study(direction="maximize")
study.optimize(objective_cv, n_trials=20, timeout=600)

PS: Unfortunatly, Optuna developers have indicated that they will remove the system_attrs in the future which I think will be a loss.