如何在Optuna交叉验证过程中记录每个折的验证损失？

我使用Toshihiko Yanase的代码在Optuna中进行超参数优化的交叉验证，代码如下：

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)

但目前该代码无法在Optuna仪表盘中记录每个折的验证损失，请问是否有方法实现这一功能？

可以通过Optuna提供的trial.report()方法，在每轮交叉验证结束后记录当前折的指标，这样就能在仪表盘中查看每个折的结果。同时还可以用trial.set_user_attr()保存每个折的具体数值，便于后续查看详情。

修改后的完整代码如下：

import optuna
import numpy as np
from sklearn.model_selection import KFold
import torch
from torch.utils.data import Subset, DataLoader
from torchvision import datasets, transforms

DIR = "./data"
BATCHSIZE = 128

def objective(trial, train_loader, valid_loader):
    # 保留原有的模型训练与验证逻辑，最终返回当前折的准确率
    # ... 此处替换为你的训练代码 ...
    return accuracy

def objective_cv(trial):
    # 获取MNIST数据集
    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 = Subset(dataset, train_idx)
        valid_data = Subset(dataset, valid_idx)

        train_loader = DataLoader(
            train_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )
        valid_loader = DataLoader(
            valid_data,
            batch_size=BATCHSIZE,
            shuffle=True,
        )

        accuracy = objective(trial, train_loader, valid_loader)
        scores.append(accuracy)
        
        # 记录当前折的准确率到Optuna，fold_idx作为步数标识
        trial.report(accuracy, fold_idx)
        # 自定义属性保存单折结果，方便后续查看
        trial.set_user_attr(f"fold_{fold_idx}_accuracy", accuracy)
        
        # 可选：如果当前trial表现太差，提前终止剪枝
        if trial.should_prune():
            raise optuna.TrialPruned()
    
    mean_accuracy = np.mean(scores)
    trial.set_user_attr("mean_accuracy", mean_accuracy)
    return mean_accuracy

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

关键说明 #

• trial.report(value, step)：将单折的准确率（或损失）与折的索引绑定，Optuna仪表盘会自动展示每折指标的变化趋势。
• trial.set_user_attr(key, value)：把每个折的具体指标值存为trial的自定义属性，在查看单个trial详情时能直接看到各折的数值。
• trial.should_prune()：可选的剪枝逻辑，若当前trial的表现远不如其他trial，可提前终止该trial，节省计算资源。

内容的提问来源于stack exchange，提问作者Gabi Gubu