线性回归任务中能否使用stratify分层抽样？有无标准实现？

回归任务里没法直接用train_test_split的stratify参数，因为这个参数要求传入离散的类别标签，而回归的目标变量是连续值。不过有个标准的变通方案——把连续目标变量离散化（分箱），之后就能用分层抽样了。

具体做法 #

核心思路是把连续的y划分成若干离散的“类别区间”，让训练集和测试集在这些区间的样本占比和原数据一致，以此保证两者的目标变量分布相似。常用的分箱方式有两种：

• 等频分箱：每个区间内的样本数量大致相等，用pd.qcut()实现，这是更推荐的方式，能避免因目标变量偏态导致某区间样本极少的问题
• 等宽分箱：把目标变量的取值范围切成宽度相同的区间，用pd.cut()实现

sklearn中的代码示例 #

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设X是特征矩阵，y是连续型目标变量
X = ...  # 替换成你的特征数据
y = ...  # 替换成你的连续目标变量

# 用等频分箱把y分成5个区间（q值可根据样本量调整，比如样本多可以设为10）
y_binned = pd.qcut(y, q=5, labels=False)

# 传入分箱后的y_binned作为stratify参数，完成分层拆分
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y_binned
)

注意事项 #

• 分箱的数量要合理：样本量小的话别分太多箱，不然某个区间可能只有几个样本，分层抽样的意义就不大了；样本量大可以适当增加箱数，让分布匹配更精准
• 如果你不想用pandas，也可以自己实现分箱逻辑，比如用numpy的percentile函数计算分位数，手动划分区间
• sklearn没有专门针对回归任务的分层抽样参数，因为分层抽样的本质是保证类别分布，所以必须先把连续目标转成离散类别才能用stratify

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