使用train_test_split分层划分数据集遇ValueError报错的解决方法

报错原因 #

使用stratify=Y进行分层抽样时，算法要求每个类别至少有2个样本（训练集和测试集需各分配至少1个），但你的标签Y中存在某个类别仅含1个样本，不符合分层抽样的最低要求，因此触发错误。

解决方案 #

• 移除样本量不足的类别
  先统计每个类别的样本数，筛选出样本量≥2的类别，过滤掉不符合要求的数据后再执行拆分：

  from collections import Counter
  # 统计各类别样本数量
  class_counts = Counter(Y)
  # 筛选有效类别（样本数≥2）
  valid_classes = [cls for cls, cnt in class_counts.items() if cnt >= 2]
  # 过滤数据与标签
  mask = [y in valid_classes for y in Y]
  X_filtered = X[mask]
  Y_filtered = Y[mask]
  # 执行分层拆分
  X_train, X_test, Y_train, Y_test = train_test_split(X_filtered, Y_filtered, test_size=0.25, stratify=Y_filtered, random_state=2)
  

• 关闭分层抽样（仅必要时使用）
  去掉stratify=Y参数，直接随机拆分数据，但这种方式可能导致训练集与测试集的类别分布和原数据不一致，仅适合样本量极大、类别分布影响可忽略的场景：

  X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.25, random_state=2)
  

• 合并小类别
  将样本量为1的稀有类别合并到相近的大类中，或者统一归为“其他”类，确保所有类别样本数≥2后再执行分层抽样。

• 过采样生成小类别样本
  使用过采样算法（如SMOTE）为样本量不足的类别生成合成样本，满足分层抽样的最低要求：

  from imblearn.over_sampling import SMOTE
  # 针对小类别执行过采样
  smote = SMOTE(random_state=2)
  X_resampled, Y_resampled = smote.fit_resample(X, Y)
  # 执行分层拆分
  X_train, X_test, Y_train, Y_test = train_test_split(X_resampled, Y_resampled, test_size=0.25, stratify=Y_resampled, random_state=2)
  

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