用鸢尾花数据集实战：5分钟搞定sklearn数据划分，附Jupyter Notebook完整代码
鸢尾花数据集实战：5分钟掌握sklearn数据划分技巧
第一次接触机器学习时，最让人头疼的往往不是算法本身，而是如何正确处理数据。记得我刚开始学习时，花了整整一个下午才搞明白怎么把数据集分成训练集和测试集。现在，让我们用最简单的鸢尾花数据集，带你快速上手sklearn的数据划分功能。
1. 环境准备与数据加载
在开始之前，确保你的Python环境中已经安装了以下库：
pip install numpy pandas scikit-learn jupyter打开Jupyter Notebook，我们首先导入必要的库并加载鸢尾花数据集：
from sklearn.datasets import load_iris import pandas as pd # 加载数据集 iris = load_iris()鸢尾花数据集是机器学习中最经典的数据集之一，包含150个样本，每个样本有4个特征：
- 萼片长度（cm）
- 萼片宽度（cm）
- 花瓣长度（cm）
- 花瓣宽度（cm）
这些特征对应三种鸢尾花的类别标签（0:山鸢尾，1:变色鸢尾，2:维吉尼亚鸢尾）。我们可以用pandas查看前5个样本：
df = pd.DataFrame(iris.data, columns=iris.feature_names) df['target'] = iris.target df.head()2. 理解数据划分的重要性
在机器学习中，我们通常不会使用全部数据来训练模型。原因很简单：我们需要一些"没见过"的数据来评估模型的真实表现。这就好比考试前，老师不会把考题直接给你背，而是留一部分题目作为真正的测试。
数据划分的基本原则：
- 训练集：用于模型训练，通常占70-80%
- 验证集：用于调参和模型选择（可选）
- 测试集：用于最终评估，通常占20-30%
注意：测试集应该只用于最终评估，不要在调参过程中反复使用，否则会导致评估结果过于乐观。
3. 使用train_test_split进行数据划分
sklearn的train_test_split函数是最常用的数据划分工具。让我们看看它的基本用法：
from sklearn.model_selection import train_test_split # 基本划分 X_train, X_test, y_train, y_test = train_test_split( iris.data, # 特征数据 iris.target, # 标签数据 test_size=0.2, # 测试集比例 random_state=42 # 随机种子 ) print(f"训练集样本数: {len(X_train)}") print(f"测试集样本数: {len(X_test)}")关键参数说明：
| 参数 | 说明 | 默认值 |
|---|---|---|
| test_size | 测试集比例 | 0.25 |
| random_state | 随机种子 | None |
| shuffle | 是否打乱数据 | True |
| stratify | 是否分层抽样 | None |
实际项目中，我强烈建议设置random_state，这样可以确保每次运行结果一致，便于调试和复现。
4. 高级划分技巧
4.1 分层抽样
当数据类别分布不均衡时，简单的随机划分可能导致某些类别在训练集或测试集中代表性不足。这时可以使用分层抽样：
X_train, X_test, y_train, y_test = train_test_split( iris.data, iris.target, test_size=0.2, stratify=iris.target, # 按标签分层 random_state=42 )4.2 多输出划分
如果你的数据有多个输出（比如多标签分类），train_test_split也能处理：
# 假设我们有两个输出 import numpy as np y2 = np.random.randint(0, 2, size=len(iris.target)) X_train, X_test, y_train, y_test, y2_train, y2_test = train_test_split( iris.data, iris.target, y2, test_size=0.2, random_state=42 )4.3 时间序列数据划分
对于时间序列数据，我们通常不希望打乱顺序：
X_train, X_test, y_train, y_test = train_test_split( iris.data, iris.target, test_size=0.2, shuffle=False # 不打乱顺序 )5. 完整示例代码
下面是一个完整的Jupyter Notebook示例，包含了数据加载、划分和简单可视化：
# 完整代码示例 from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import pandas as pd import matplotlib.pyplot as plt # 加载数据 iris = load_iris() df = pd.DataFrame(iris.data, columns=iris.feature_names) df['target'] = iris.target # 数据划分 X_train, X_test, y_train, y_test = train_test_split( iris.data, iris.target, test_size=0.2, stratify=iris.target, random_state=42 ) # 可视化类别分布 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) df['target'].value_counts().plot(kind='bar', ax=ax1, title='原始数据分布') pd.Series(y_train).value_counts().plot(kind='bar', ax=ax2, title='训练集分布') plt.show() print("训练集形状:", X_train.shape) print("测试集形状:", X_test.shape)运行这段代码，你会看到原始数据和训练集的类别分布对比，确保划分后的数据保持了原始分布。
6. 常见问题与解决方案
在实际项目中，数据划分可能会遇到各种问题。以下是我总结的一些常见情况及解决方法：
问题1：数据集太小，划分后训练样本不足
解决方案：
- 使用交叉验证代替简单划分
- 尝试数据增强技术
- 调整test_size参数，减少测试集比例
问题2：类别极度不均衡
解决方案：
- 使用分层抽样（stratify参数）
- 考虑过采样或欠采样技术
- 使用类别权重参数
问题3：特征和标签不在同一个DataFrame中
解决方案：
# 假设features和labels是两个独立的DataFrame X_train, X_test, y_train, y_test = train_test_split( features.values, # 转换为numpy数组 labels.values.ravel(), # 确保是一维数组 test_size=0.2 )问题4：需要划分多个数据集
解决方案：
# 划分训练、验证、测试集 X_train, X_temp, y_train, y_temp = train_test_split( X, y, test_size=0.4, random_state=42 ) X_val, X_test, y_val, y_test = train_test_split( X_temp, y_temp, test_size=0.5, random_state=42 )7. 实际应用建议
经过多个项目的实践，我发现这些技巧特别有用：
- 保持一致性：在整个项目中固定random_state值，确保每次运行结果一致
- 尽早划分：在数据探索前就划分好测试集，避免数据泄露
- 考虑业务场景：某些场景可能需要按时间划分（如预测未来数据）
- 文档记录：记录划分比例、随机种子等参数，便于复现
对于初学者，我建议先从简单的划分开始，等熟悉基本流程后，再尝试更复杂的交叉验证方法。记住，数据划分是机器学习工作流中最基础的步骤之一，但也是最容易出错的地方。