PyTorch分层划分训练/验证/测试集时第六类全入测试集问题排查

核心错误分析 #

1. 索引误用（最关键问题） #

原代码第二次划分时犯了索引匹配错误：

• 第一次划分后，train_dataset是基于原始dataset的Subset，它的内部索引是相对索引（从0到len(train_dataset)-1），对应原始dataset的train_indices列表中的元素。
• 第二次train_test_split得到的train_indices是这个train_dataset的相对索引，但你直接用它去创建Subset(dataset, train_indices)，这相当于取原始dataset的第0、1、2...个样本，完全偏离了第一次划分后的结果，直接破坏了分层逻辑，甚至会导致某些类的样本被完全排除出训练集，看起来像是全进入了测试集。

2. 小类别样本量导致的分层异常 #

如果第六类的总样本数极少（比如≤10），第一次用test_size=0.1做分层划分时，sklearn的train_test_split无法按比例拆分（不可能取0.1个样本），会直接将该类的全部样本分配到测试集，这是分层划分在样本量不足时的必然结果。

修正后的代码 #

方案一：基于Subset嵌套划分（推荐） #

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

# 第一步：划分训练+验证集（90%）与测试集（10%）
train_val_indices, test_indices, _, _ = train_test_split(
    range(len(dataset)),
    dataset.targets,
    stratify=dataset.targets,
    test_size=0.1,
    random_state=1
)

train_val_dataset = Subset(dataset, train_val_indices)
test_dataset = Subset(dataset, test_indices)

# 第二步：从训练+验证集中拆分训练集（80%原始占比）与验证集（10%原始占比）
train_indices, val_indices, _, _ = train_test_split(
    range(len(train_val_dataset)),
    # 直接取原始数据集的对应标签，避免生成器遍历问题
    [dataset.targets[idx] for idx in train_val_indices],
    stratify=[dataset.targets[idx] for idx in train_val_indices],
    test_size=0.111,
    random_state=1
)

# 基于train_val_dataset创建最终的训练/验证集，保持索引匹配
train_dataset = Subset(train_val_dataset, train_indices)
val_dataset = Subset(train_val_dataset, val_indices)

方案二：转换为原始数据集索引 #

如果需要直接基于原始dataset创建所有子集，可以将相对索引转换为原始索引：

# 第二步替换为：
train_indices, val_indices, _, _ = train_test_split(
    range(len(train_val_dataset)),
    [dataset.targets[idx] for idx in train_val_indices],
    stratify=[dataset.targets[idx] for idx in train_val_indices],
    test_size=0.111,
    random_state=1
)

# 将相对索引转换为原始数据集的索引
train_indices_original = [train_val_indices[i] for i in train_indices]
val_indices_original = [train_val_indices[i] for i in val_indices]

train_dataset = Subset(dataset, train_indices_original)
val_dataset = Subset(dataset, val_indices_original)

小类别样本的处理建议 #

如果第六类样本量极少，可采取以下方式优化：

• 调整划分比例：比如先确保每个类在测试集中至少有1个样本，再划分剩余样本。
• 使用分层K折交叉验证：比如采用10折，将其中1折作为测试集，另外9折中再拆分训练/验证集，保证小类别在所有集合中都有分布。
• 数据增强：对小类别样本进行数据增强，增加样本数量后再划分。

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