PyTorch中K折验证测试阶段Dataloader索引错误求解

问题场景 #

使用sklearn的KFold模块进行分类器的K折训练，训练过程正常，但测试阶段触发IndexError，核心错误信息如下：

IndexError: index 106 is out of bounds for dimension 0 with size 27

错误原因分析 #

1. 验证集数据集不匹配：KFold是基于torch_trainDataset拆分的索引，但验证集DataLoader错误使用了torch_testDataset，测试集样本量远小于训练集，导致训练集的索引在测试集里不存在，直接触发索引越界。
2. test函数循环变量错误：test函数内部循环遍历的是未定义的data_train，而非传入的data_test参数。
3. train函数累加逻辑错误：train函数将外部的train_loss作为参数传入并再次累加，会导致损失值重复计算，数值异常。

修复方案 #

• 验证集必须使用训练集的Dataset，因为KFold生成的索引仅适用于训练集。
• 修正test函数的循环变量，改为遍历传入的data_test。
• 重构train函数的损失累加逻辑，移除外部传入的train_loss参数，在函数内部初始化局部变量进行累加。
• 验证准确率计算时，基于验证集的实际样本数（即len(valid_idx)），而非DataLoader的长度。

完整修正代码 #

K折训练循环 #

from sklearn.model_selection import KFold
from torch.utils.data import SubsetRandomSampler, DataLoader

FK_split = KFold(n_splits=4, shuffle=True, random_state=0)

for epoch in range(num_epoch):
    train_loss = 0.0
    Acc_valid = 0.0

    for train_idx, valid_idx in FK_split.split(torch_trainDataset):
        train_sampler = SubsetRandomSampler(train_idx)
        valid_sampler = SubsetRandomSampler(valid_idx)
        # 验证集使用训练集的Dataset，而非测试集
        train_dataloader = DataLoader(torch_trainDataset, batch_size=1, sampler=train_sampler)
        valid_dataloader = DataLoader(torch_trainDataset, batch_size=1, sampler=valid_sampler)

        # 调用修正后的train函数
        fold_train_loss = train(model, train_dataloader, lossfunc, optimizer)
        train_loss += fold_train_loss
        # 调用修正后的test函数
        _, acc_valid = test(model, valid_dataloader)
        Acc_valid += acc_valid

    # 计算平均训练损失和验证准确率
    avg_train_loss = train_loss / len(torch_trainDataset)
    avg_valid_acc = Acc_valid / len(torch_trainDataset)
    print(f"Epoch {epoch+1}: Train Loss = {avg_train_loss:.4f}, Valid Acc = {avg_valid_acc:.4f}")

修正后的train、test、get_acc函数 #

def train(model, data_train, lossfunc, optimizer):
    model.train()
    train_loss = 0.0
    for x, y in data_train:
        optimizer.zero_grad()
        output = model(x)
        loss = lossfunc(output, y)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * x.size(0)
    return train_loss

def get_acc(outputs, labels):
    """计算准确率"""
    _, predict = torch.max(outputs.data, 1)
    correct_num = (labels == predict).sum().item()
    return predict, correct_num

def test(model, data_test):
    model.eval()
    Predict = []
    Acc = 0.0
    # 修正循环变量为传入的data_test
    with torch.no_grad():
        for x, y in data_test:
            outputs = model(x)
            predict, acc = get_acc(outputs, y)
            Predict.append(predict.tolist())
            Acc += acc
    return Predict, Acc

额外说明 #

• 测试集应该在所有K折训练完成后单独使用，用来评估模型的最终泛化能力，而不是在每折训练中作为验证集。
• 在test函数中添加model.eval()和torch.no_grad()，可以关闭模型的训练模式，避免梯度计算，提升测试效率并节省内存。

内容的提问来源于stack exchange，提问作者Ryan丶