企业官网建设流程全解析

鸢尾花数据集实战：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个特征：

1. 萼片长度（cm）
2. 萼片宽度（cm）
3. 花瓣长度（cm）
4. 花瓣宽度（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)}")

关键参数说明：

实际项目中，我强烈建议设置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. 实际应用建议

经过多个项目的实践，我发现这些技巧特别有用：

1. 保持一致性：在整个项目中固定random_state值，确保每次运行结果一致
2. 尽早划分：在数据探索前就划分好测试集，避免数据泄露
3. 考虑业务场景：某些场景可能需要按时间划分（如预测未来数据）
4. 文档记录：记录划分比例、随机种子等参数，便于复现

对于初学者，我建议先从简单的划分开始，等熟悉基本流程后，再尝试更复杂的交叉验证方法。记住，数据划分是机器学习工作流中最基础的步骤之一，但也是最容易出错的地方。