如何用pandas和sklearn将数据集按正负样本比例8:2划分训练集和测试集？

Question

如何用Pandas和Sklearn按正负样本比例8:2划分数据集？ 在处理分类问题时，确保训练集和测试集的正负样本比例一致至关重要。使用Pandas加载和预处理数据后，如何借助Sklearn的`train_test_split`方法实现按8:2比例划分，并保持正负样本分布均衡？需设置`stratify`参数为标签列，同时指定`test_size=0.2`。例如：`X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)`。但若数据集中存在严重类别不均衡，该如何进一步优化划分策略以保证模型训练效果？

Answer 1

1. 数据集划分基础

在机器学习中，数据集的划分是构建模型的关键步骤之一。为了确保训练集和测试集具有相同的正负样本分布，我们需要使用`train_test_split`方法中的`stratify`参数。以下是一个简单的示例代码：

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设数据已加载到DataFrame中
data = pd.read_csv('data.csv')
X = data.drop('target', axis=1)
y = data['target']

# 按8:2比例划分数据集，并保持正负样本分布均衡
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)
    

上述代码中，`test_size=0.2`表示测试集占20%，`stratify=y`确保正负样本比例一致。

2. 类别不均衡问题分析

当数据集中存在严重类别不均衡时，即使使用`stratify`参数，也可能无法充分优化模型效果。以下是常见的类别不均衡场景及影响：

• 正样本占比远低于负样本（如欺诈检测、疾病诊断等）。
• 少数类样本可能被忽略或过拟合。
• 模型性能指标（如准确率）可能失真。

针对这些问题，我们需要进一步优化划分策略。

3. 优化策略：重采样与分层抽样结合

为了解决类别不均衡问题，可以采用以下方法：

方法	描述	适用场景
欠采样	减少多数类样本数量以平衡数据集。	多数类数据充足且不影响模型泛化能力。
过采样	增加少数类样本数量，例如使用SMOTE算法生成合成样本。	少数类数据不足但需保留更多信息。
组合采样	同时应用欠采样和过采样技术。	数据集既需要减少噪声又需要增强少数类。

结合分层抽样，可以在划分数据集之前先进行重采样处理。

4. 流程图：完整数据处理流程

以下是完整的数据处理流程，包含数据加载、预处理、重采样和分层划分步骤：

graph TD;
            A[加载数据] --> B[检查类别分布];
            B --> C{类别是否均衡?};
            C --否--> D[应用重采样技术];
            D --> E[划分训练集和测试集];
            C --是--> E;
            E --> F[训练模型];

通过该流程，我们可以确保数据集划分合理并有效应对类别不均衡问题。

5. 示例代码：结合SMOTE与分层划分

以下代码展示了如何结合SMOTE算法和`train_test_split`实现优化划分：

from imblearn.over_sampling import SMOTE

# 应用SMOTE算法进行过采样
smote = SMOTE(random_state=42)
X_resampled, y_resampled = smote.fit_resample(X, y)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X_resampled, y_resampled, test_size=0.2, stratify=y_resampled, random_state=42
)
    

此代码首先使用SMOTE算法对少数类进行过采样，然后通过`train_test_split`按8:2比例划分数据集。