拆分大型文本数据集用于数据匹配应采用何种sklearn分层抽样技术？

sklearn中完全支持你需要的这种分层拆分能力，对应工具为StratifiedKFold：

• 该类属于sklearn.model_selection模块，核心能力就是将数据集拆分为N个互斥子集，拆分过程会严格保证每个子集的类别占比和原数据集完全一致，完美适配二分类、多分类场景的分层拆分需求。
• 你不需要用到它的交叉验证能力，只需要遍历拆分结果取出每个子集的索引，即可拿到所有符合要求的独立子集。

示例代码 #

from sklearn.model_selection import StratifiedKFold
import numpy as np

# 替换为你的实际特征矩阵、标签数组
X = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
y = np.array([0, 0, 1, 1, 0, 1])

# 定义要拆分的子集数量N
n_subsets = 3
# shuffle=True会先打乱数据集再拆分，random_state固定结果方便复现
skf = StratifiedKFold(n_splits=n_subsets, shuffle=True, random_state=42)

# 存储所有分层子集
subset_list = []
for _, subset_idx in skf.split(X, y):
    X_sub = X[subset_idx]
    y_sub = y[subset_idx]
    subset_list.append((X_sub, y_sub))

注意事项 #

• 你已经完成了类别不平衡处理，直接用处理后的标签y作为分层依据即可，不需要额外调整参数。
• 只要原数据集总样本数可以被N整除，拆分出来的每个子集样本量完全相等；如果不能整除，子集样本量最多只有1个的差异，不会影响类别占比。
• 拿到所有子集后，你可以对每个子集单独调用train_test_split拆分训练集和测试集，独立拟合分类器即可。

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