如何对高度不平衡二分类数据集进行自定义采样以划分训练测试集？

嗨，针对你的高度不平衡数据集问题，我来给你梳理几种实用的自定义采样实现方式，都是在你现有train_test_split代码基础上扩展的，而且要记住：采样只应该在训练集上进行，绝对不能碰测试集，避免数据泄露影响模型评估的准确性。

欠采样是从数量多的标签0样本里挑选一部分，让它和标签1的样本数量接近。

1. 随机欠采样（简单直接） #

你可以用imblearn库的工具，也可以手动实现：

方法一：用imblearn库 #

先安装依赖（如果没装的话）：

pip install imblearn

然后修改代码（注意：sklearn.cross_validation已经废弃，建议替换为sklearn.model_selection）：

from sklearn.model_selection import train_test_split
from imblearn.under_sampling import RandomUnderSampler

X = Actual_DataFrame.copy()
y = X.pop('Attrition')

# 先划分训练测试集，保持分层抽样
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.20, random_state=42, stratify=y)

# 初始化欠采样器：sampling_strategy设为'auto'会自动让两类样本数量一致
rus = RandomUnderSampler(sampling_strategy='auto', random_state=42)
X_train_resampled, y_train_resampled = rus.fit_resample(X_train, y_train)

# 现在X_train_resampled和y_train_resampled就是平衡后的训练集了

方法二：手动实现随机欠采样 #

如果不想用第三方库，自己写也很简单：

import pandas as pd

# 划分训练测试集后，合并成完整训练DataFrame
train_df = pd.concat([X_train, y_train], axis=1)

# 分离多数类（标签0）和少数类（标签1）
class_0 = train_df[train_df['Attrition'] == 0]
class_1 = train_df[train_df['Attrition'] == 1]

# 从多数类里随机挑选和少数类数量相等的样本
class_0_sampled = class_0.sample(n=len(class_1), random_state=42)

# 合并成平衡训练集
balanced_train_df = pd.concat([class_0_sampled, class_1], axis=0)

# 拆分回特征和标签
X_train_resampled = balanced_train_df.drop('Attrition', axis=1)
y_train_resampled = balanced_train_df['Attrition']

过采样是对数量少的标签1样本进行复制或生成合成样本，让它和标签0的数量匹配。

1. 随机过采样（简单复制） #

手动实现逻辑如下：

import pandas as pd

# 划分训练测试集后合并训练集
train_df = pd.concat([X_train, y_train], axis=1)

class_0 = train_df[train_df['Attrition'] == 0]
class_1 = train_df[train_df['Attrition'] == 1]

# 重复采样少数类样本，直到和多数类数量一致（replace=True允许重复采样）
class_1_sampled = class_1.sample(n=len(class_0), replace=True, random_state=42)

# 合并成平衡训练集
balanced_train_df = pd.concat([class_0, class_1_sampled], axis=0)

X_train_resampled = balanced_train_df.drop('Attrition', axis=1)
y_train_resampled = balanced_train_df['Attrition']

2. SMOTE合成过采样（生成新样本） #

SMOTE通过插值生成新的少数类样本，比单纯复制更能避免过拟合，同样用imblearn实现：

from imblearn.over_sampling import SMOTE

# 划分训练测试集后
smote = SMOTE(sampling_strategy='auto', random_state=42)
X_train_resampled, y_train_resampled = smote.fit_resample(X_train, y_train)

单独欠采样可能丢失多数类信息，单独过采样容易过拟合，混合采样是折中方案，比如先用SMOTE过采样少数类，再对多数类做欠采样：

from imblearn.combine import SMOTEENN

smote_enn = SMOTEENN(random_state=42)
X_train_resampled, y_train_resampled = smote_enn.fit_resample(X_train, y_train)

如果你想根据特定规则采样（比如选择多数类中特征分布和少数类接近的样本），可以自己写逻辑，比如基于KNN选择近邻多数类样本：

from sklearn.neighbors import NearestNeighbors
import numpy as np
import pandas as pd

# 划分训练测试集后
train_df = pd.concat([X_train, y_train], axis=1)
class_0 = train_df[train_df['Attrition'] == 0]
class_1 = train_df[train_df['Attrition'] == 1]

# 用少数类样本训练KNN，找到每个少数类样本的5个近邻多数类样本
nn = NearestNeighbors(n_neighbors=5)
nn.fit(class_0.drop('Attrition', axis=1))
_, indices = nn.kneighbors(class_1.drop('Attrition', axis=1))

# 收集这些近邻样本并去重
selected_class_0_indices = np.unique(indices.flatten())
class_0_sampled = class_0.iloc[selected_class_0_indices]

# 合并成平衡训练集
balanced_train_df = pd.concat([class_0_sampled, class_1], axis=0)
X_train_resampled = balanced_train_df.drop('Attrition', axis=1)
y_train_resampled = balanced_train_df['Attrition']

最后提醒：采样后训练模型时，评估一定要用原始的X_test和y_test，这样才能真实反映模型在不平衡数据上的泛化能力。

内容的提问来源于stack exchange，提问作者Student of the Digital World