I think we need to evaluate all folds and calculate the mean inside an objective function. I create an example notebook, so please take a look.
In the notebook, I slightly modified the objective
function to pass the dataset with the arguments and added a wrapper function objective_cv
to call the objective
function with the split dataset. Then, I optimized the objective_cv
instead of the objective
function.
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)