如何让sklearn train_test_split基于自定义二元向量prop实现分层划分？

我想用sklearn的train_test_split函数将数据划分为训练集、验证集和测试集，指定自定义二元向量prop作为stratify参数，让划分后各数据集的标签比例与prop的比例一致，但当前代码未达到预期效果。

原代码 #

import numpy as np
from sklearn.model_selection import train_test_split

# 假设y为原始标签，self.train_data等为自定义数据结构
prop = np.zeros(len(y))
one_amount = round(0.012 * len(y))
prop[:one_amount] = 1

self.train_data.X, self.test_data.X, self.train_data.y, self.test_data.y = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=prop
)

self.train_data.X, self.val_data.X, self.train_data.y, self.val_data.y = train_test_split(
    self.train_data.X, self.train_data.y, test_size=0.25, random_state=42
)

划分结果 #

生成的prop向量及各数据集标签比例：

• prop标签比例：class 0: 10988 (99.0%), class 1: 133 (1.0%)
• 训练集标签比例：class 0: 4989 (75.0%), class 1: 1683 (25.0%)
• 验证集标签比例：class 0: 1665 (75.0%), class 1: 559 (25.0%)
• 测试集标签比例：class 0: 1676 (75.0%), class 1: 549 (25.0%)

问题原因 #

核心问题出在第二次划分：第二次调用train_test_split时没有传入stratify参数，且第一次划分用prop作为分层依据，但第二次直接基于原始标签y随机划分，完全丢失了prop的分层逻辑，最终导致各数据集比例偏离预期。

修正方案 #

要全程保持分层逻辑一致，两次划分都要基于prop向量进行分层：

修改后的代码 #

import numpy as np
from sklearn.model_selection import train_test_split

# 生成自定义分层向量prop
prop = np.zeros(len(y))
one_amount = round(0.012 * len(y))
prop[:one_amount] = 1

# 第一步：划分「训练+验证集」和「测试集」，同步拆分prop以保留分层依据
X_train_val, X_test, y_train_val, y_test, prop_train_val, _ = train_test_split(
    X, y, prop, test_size=0.2, random_state=42, stratify=prop
)

# 第二步：从「训练+验证集」中划分「训练集」和「验证集」，用对应部分的prop分层
X_train, X_val, y_train, y_val = train_test_split(
    X_train_val, y_train_val, test_size=0.25, random_state=42, stratify=prop_train_val
)

# 赋值到自定义数据结构
self.train_data.X, self.train_data.y = X_train, y_train
self.val_data.X, self.val_data.y = X_val, y_val
self.test_data.X, self.test_data.y = X_test, y_test

关键说明 #

• 第一次划分时同步拆分prop，得到对应「训练+验证集」的prop_train_val，确保第二次划分时能沿用相同的分层基准。
• 全程以prop作为stratify参数的输入，才能保证各数据集的标签比例和prop的比例一致。

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