如何指定类别占比划分训练测试集使测试集正样本占比约50%

sklearn.model_selection.train_test_split的stratify参数设计目标是让划分后的训练集、测试集类别占比和原始数据集完全一致，原始数据正样本（y=1）占比低于50%的前提下，用这个参数不可能得到正样本占比50%的测试集，不是调用方式错了，是参数本身不支持自定义测试集的类别占比。

最直接的实现逻辑是按目标占比分别从两个类别中独立抽样组成测试集，剩余样本全部划入训练集，全程不要对测试集样本做任何生成、复制类的操作，避免评估指标失真。

具体步骤 #

• 先把全量数据集按标签拆成正样本（y=1）、负样本（y=0）两个独立子集
• 确定测试集的可承载规模：受限于正样本总数，测试集最大规模为「正样本总数*2」（此时测试集用完全部正样本，再抽同等数量的负样本配对，刚好1:1）；如果不想用完全部正样本，也可以按训练需求缩减测试集规模，只要保证抽取的正、负样本数相等即可
• 分别从两个子集里无放回随机抽取对应数量的样本，拼接后打乱就是正样本占比50%的测试集
• 两个子集抽剩的所有样本直接拼接打乱，就是训练集。训练集保留原始类别分布即可，后续训练阶段如果要处理类别不平衡，可以单独对训练集做重采样、加类别权重，不要提前改动测试集

代码参考 #

import pandas as pd
from sklearn.utils import shuffle

# 假设总数据集存储在df中，标签列名为'y'
pos_samples = df[df["y"] == 1].reset_index(drop=True)
neg_samples = df[df["y"] == 0].reset_index(drop=True)

# 配置测试集抽样数：例如下方取80%的正样本进测试集，抽同等数量负样本配对
# 如果要用完所有正样本，把test_pos_count改成len(pos_samples)即可
test_pos_count = int(0.8 * len(pos_samples))
test_neg_count = test_pos_count  # 保证测试集两类样本数相等，占比各50%

# 抽样生成测试集
test_pos = pos_samples.sample(n=test_pos_count, random_state=42)
test_neg = neg_samples.sample(n=test_neg_count, random_state=42)
test_set = pd.concat([test_pos, test_neg]).pipe(shuffle, random_state=42)

# 剩余样本生成训练集
train_pos = pos_samples.drop(test_pos.index)
train_neg = neg_samples.drop(test_neg.index)
train_set = pd.concat([train_pos, train_neg]).pipe(shuffle, random_state=42)

# 拆分特征和标签
X_train, y_train = train_set.drop(columns=["y"]), train_set["y"]
X_test, y_test = test_set.drop(columns=["y"]), test_set["y"]

• 绝对不要对测试集做任何过采样操作（比如复制正样本、用SMOTE等算法生成合成正样本），测试集必须全部由真实样本组成，否则算出来的模型准确率、召回率等指标没有任何实际参考价值
• 不要为了凑测试集占比把训练集的正样本抽得太少，至少要保证训练集里留存足够的正样本供模型学习，否则模型效果会大打折扣
• 如果不需要强制让y_test的物理占比为50%，只是想公平评估不平衡场景下的模型效果，也可以保持分层抽样的测试集原始分布，评估时用F1、PR-AUC、混淆矩阵加权指标等对类别不平衡不敏感的评估方式即可

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