Pytorch 中的分层训练/验证/测试拆分

Question

我有一个包含 6 个类别的图像分类数据集，我正在使用 torchvision ImageFolder 类加载该数据集。我编写了以下内容，以分层方式将数据集分为 3 组：

from torch.utils.data import Subset
from sklearn.model_selection import train_test_split

train_indices, test_indices, _, _ = train_test_split
(
    range(len(dataset)),
    dataset.targets,
    stratify=dataset.targets,
    test_size=0.1,
    random_state=1
)

train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

train_targets = [label for _, label in train_dataset]
train_indices, val_indices, _, _ = train_test_split
(
    range(len(train_dataset)),
    train_targets,
    stratify=train_targets,
    test_size=0.111,
    random_state=1
)

train_dataset = Subset(dataset, train_indices)

这可以正确地将数据分割为 90/10/10%。除第六级外，所有班级在所有三组之间都遵循相同的分布。属于第 6 类的所有样本最终都会进入测试集。

我做错了什么？

谢谢。

Answer 1

我也遇到了同样的问题。答案是您的数据集和子集不匹配。 当您构建第二个分割时，您使用 train_dataset 来构建索引，然后最后使用 dataset 作为 train_dataset 子集。这是一个错误，当您在其上构建索引时，它应该是第一次分割时的 train_dataset，而不是在 dataset 上。

我也做了一些优化，以免在第一次拆分时重建标签。 应该是：

train_indices, test_indices, train_label, _ = train_test_split
(
    range(len(dataset)),
    dataset.targets,
    stratify=dataset.targets,
    test_size=0.1,
    random_state=1
)

train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

train_indices, val_indices, _, _ = train_test_split
(
    range(len(train_dataset)),
    train_label,
    stratify=train_label,
    test_size=0.111,
    random_state=1
)

train_dataset = Subset(train_dataset, train_indices)
val_dataset = Subset(train_dataset, val_indices)