Sklearn中如何实现训练/测试/验证集的分层拆分

完全可以基于scikit-learn实现符合要求的分层拆分，无需额外依赖其他第三方库，核心思路是分两次调用train_test_split接口，两次拆分均开启分层抽样参数，即可保证训练集、验证集、测试集的类别分布与原数据集完全一致。

拆分比例计算逻辑 #

两次拆分的参数需要提前算对，避免最终比例不符合要求：

• 第一次拆分：从全量数据中分出70%作为训练集，剩余30%作为临时集合（用于后续拆分验证集、测试集），第一次拆分的test_size参数设为0.3即可。
• 第二次拆分：对占原数据30%的临时集合做二次拆分，最终需要得到占原数据20%的验证集、占原数据10%的测试集，因此临时集合中验证集占比为2/3、测试集占比为1/3，第二次拆分的test_size参数设为1/3即可。

完整代码实现 #

直接拆分pandas DataFrame（适配DataFrame存储全量数据的场景） #

假设全量数据集存在名为df的pandas DataFrame中，存储类别标签的列名为label，拆分代码如下：

import pandas as pd
from sklearn.model_selection import train_test_split

# 第一次拆分：分出70%训练集，剩余30%为临时集合
df_train, df_temp = train_test_split(
    df,
    test_size=0.3,
    stratify=df['label'],  # 指定按label列做分层抽样
    random_state=42  # 固定随机种子保证结果可复现
)

# 第二次拆分：临时集合拆分为20%原数据量的验证集、10%原数据量的测试集
df_val, df_test = train_test_split(
    df_temp,
    test_size=1/3,
    stratify=df_temp['label'],  # 临时集同样按label列分层
    random_state=42
)

拆分提前分离的特征、标签对象 #

如果已经提前把特征矩阵、标签列拆分为X（特征）、y（标签）两个对象，可以用如下写法，逻辑完全一致：

# 分离特征和标签
X = df.drop('label', axis=1)
y = df['label']

# 第一次拆分
X_train, X_temp, y_train, y_temp = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=42
)

# 第二次拆分
X_val, X_test, y_val, y_test = train_test_split(
    X_temp, y_temp, test_size=1/3, stratify=y_temp, random_state=42
)

拆分结果校验 #

拆分完成后可以运行如下代码，校验样本占比和类别分布是否符合要求：

# 校验各子集样本量占比
print(f"训练集样本占比：{len(df_train)/len(df):.2f}")
print(f"验证集样本占比：{len(df_val)/len(df):.2f}")
print(f"测试集样本占比：{len(df_test)/len(df):.2f}")

# 校验各子集类别分布
print("\n原数据集类别占比：")
print(df['label'].value_counts(normalize=True).round(4))
print("\n训练集类别占比：")
print(df_train['label'].value_counts(normalize=True).round(4))
print("\n验证集类别占比：")
print(df_val['label'].value_counts(normalize=True).round(4))
print("\n测试集类别占比：")
print(df_test['label'].value_counts(normalize=True).round(4))

正常情况下运行后可以看到三个子集的样本占比分别为0.7、0.2、0.1，且各类别的占比和原数据集几乎完全一致。

注意事项 #

• 两次拆分都必须传入stratify参数指定分层依据的标签列，否则第二次拆分无法保证验证集、测试集的类别分布和原数据对齐。
• 两次拆分建议固定相同的random_state值，保证每次运行得到的拆分结果一致，方便后续模型实验复现。
• 该方法对样本量极少的类别同样生效，分层抽样逻辑会自动按比例分配每个类别的样本到三个子集中，不会出现小样本类别在某一子集中完全缺失的问题。

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