学习笔记(29):训练集与测试集划分详解：train_test_split 函数深度解析

学习笔记(29):训练集与测试集划分详解：train_test_split 函数深度解析

一、为什么需要划分训练集和测试集？

在机器学习中，模型需要经历两个核心阶段：

1. 训练阶段：用训练集数据学习特征与目标值的映射关系（如线性回归的权重）。
2. 测试阶段：用测试集评估模型在未见过的数据上的表现，避免 “过拟合”（模型只记住训练数据的噪声，无法泛化到新数据）。

类比场景：学生通过 “练习题”（训练集）学习知识，再通过 “考试题”（测试集）检验真实水平。

二、train_test_split 函数的核心参数与逻辑

X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=42
)

1. 输入参数解析

• X_scaled：特征矩阵（已标准化的面积、房龄等特征）。
• y：目标变量（房价）。
• test_size=0.2：测试集占总数据的比例（20%），也可设为整数（如 test_size=20 表示取 20 个样本）。
• random_state=42：随机种子，确保每次划分结果一致（与 np.random.seed(42) 作用类似）。

2. 划分逻辑

• 随机抽样：按 test_size 比例从原始数据中随机抽取样本作为测试集，剩余作为训练集。
• 数据对齐：确保 X 和 y 的样本顺序一一对应（如第 i 个特征向量对应第 i 个房价标签）。

三、划分结果的维度与含义

假设原始数据有 100 个样本（n_samples=100）：

• 训练集：80 个样本（X_train.shape=(80, 2), y_train.shape=(80,)），用于模型学习。
• 测试集：20 个样本（X_test.shape=(20, 2), y_test.shape=(20,)），用于评估模型泛化能力。

四、关键参数深度解析

1. test_size：平衡训练与测试的样本量

• 取值建议：
  • 小数据集（<1000 样本）：常用 test_size=0.2~0.3（20%-30% 作为测试集）。
  • 大数据集（>10000 样本）：可设 test_size=0.1 甚至更低（因少量样本已足够评估）。
• 极端案例：若 test_size=1.0，则所有数据都是测试集，无训练集；若 test_size=0，则全是训练集。

2. random_state：确保可复现的 “随机” 划分

• 作用：固定随机种子后，每次运行代码时，训练集和测试集的样本索引完全相同。
• 示例对比：
  • 不设置 random_state：每次划分结果不同，导致模型评估指标波动。
  • 设置 random_state=42：多次运行代码，划分结果一致，便于对比不同模型效果。

3. shuffle=True（默认参数）：打乱数据顺序

• 为什么需要打乱？
  若数据按顺序排列（如前 50 个是小户型，后 50 个是大户型），不打乱会导致训练集和测试集样本分布不均（如测试集全是大户型）。
• 参数设置：train_test_split 默认为 shuffle=True，即先打乱数据再划分；若数据已随机排列，可设 shuffle=False。

五、进阶应用：分层抽样（Stratified Sampling）

当目标变量是分类变量（如二分类 “是否违约”）时，普通随机划分可能导致训练 / 测试集的类别比例失衡（如测试集全是 “违约” 样本）。此时需用 StratifiedShuffleSplit 实现分层抽样：

from sklearn.model_selection import StratifiedShuffleSplit

# 4. 使用分层抽样（确保类别比例平衡）
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_idx, test_idx in sss.split(X_scaled, y_binary):
    X_train, X_test = X_scaled[train_idx], X_scaled[test_idx]
    y_train, y_test = y_binary[train_idx], y_binary[test_idx]

print("===== 分类模型结果 =====")
print(f"原始数据类别比例：{np.bincount(y_binary)/len(y_binary)}")
print(f"训练集类别比例：{np.bincount(y_train)/len(y_train)}")
print(f"测试集类别比例：{np.bincount(y_test)/len(y_test)}")

六、实战误区与注意事项

1. 禁止在测试集上训练：测试集只能用于评估，若根据测试集结果调整模型参数（如调优正则化系数），本质上是 “偷看答案”，会导致评估结果过于乐观。
2. 数据标准化的顺序：
   • 正确流程：先划分训练测试集，再对训练集拟合标准化器（scaler.fit(X_train)），最后用训练集的标准化参数转换测试集（scaler.transform(X_test)）。
   • 错误操作：对全量数据标准化后再划分，会导致测试集 “偷看到” 全量数据的统计特征，违反 “未知数据” 假设。
3. 多轮划分与交叉验证：当数据量较小时，可使用 K 折交叉验证（如 10 折），将数据分成 10 份，每次用 9 份训练、1 份测试，重复 10 次取平均，减少单次划分的随机性误差。

七、总结：划分训练测试集的核心原则

1. 独立性：测试集数据必须是模型未见过的，模拟真实应用场景。
2. 代表性：训练集和测试集的样本分布应尽可能一致（如特征取值范围、类别比例）。
3. 可复现性：通过设置随机种子，确保实验结果可重复验证。

通过合理划分训练集与测试集，你可以更准确地评估模型的实际能力，避免被 “过拟合” 的假象误导 —— 这是机器学习工程化中至关重要的一步！

二分类问题（房价是否高于中位数）-全代码

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split, StratifiedShuffleSplit
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score

# 配置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 生成模拟数据（假设房价与面积、房龄的关系）
np.random.seed(42)
n_samples = 100
# 面积（平方米），房龄（年）
X = np.random.rand(n_samples, 2) * 100
X[:, 0] = X[:, 0]  # 面积范围：0-100
X[:, 1] = X[:, 1]  # 房龄范围：0-100

# 真实房价 = 5000*面积 + 1000*房龄 + 随机噪声（模拟真实场景）
y = 5000 * X[:, 0] + 1000 * X[:, 1] + np.random.randn(n_samples) * 10000

# 2. 将连续的房价y转换为分类标签（例如分为低、中、高3个类别）
y_category = pd.qcut(y, q=3, labels=[0, 1, 2])  # 使用pandas的qcut进行分位数切割

# 3. 数据预处理：标准化特征
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 4. 使用分层抽样
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_idx, test_idx in sss.split(X_scaled, y_category):
    X_train, X_test = X_scaled[train_idx], X_scaled[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]  # 注意：这里仍然使用原始的连续房价作为目标_

# 确保训练集和测试集的类别比例与原始数据一致
print(f"原始数据类别比例：{np.bincount(y_category)/len(y_category)}")
print(f"训练集类别比例：{np.bincount(y_category[train_idx])/len(y_category[train_idx])}")
print(f"测试集类别比例：{np.bincount(y_category[test_idx])/len(y_category[test_idx])}")

# 后续回归模型训练和评估代码保持不变
model = LinearRegression()
model.fit(X_train, y_train)
y_pred = model.predict(X_test)

# 评估模型
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print(f"均方误差: {mse:.2f}")
print(f"决定系数R²: {r2:.2f}")

打印：

原始数据类别比例：[0.34 0.32 0.34]
训练集类别比例：[0.3375 0.325  0.3375]
测试集类别比例：[0.35 0.3  0.35]
均方误差: 101112597.45
决定系数R²: 1.00

代码解析：

核心步骤解析

1. 数据准备与二分类转换

   • 生成与方案 1 相同的模拟数据（面积、房龄 → 房价）。
   • 将连续的房价y转换为二分类标签：

threshold = np.median(y)  # 使用中位数作为阈值
y_binary = (y > threshold).astype(int)  # 0=低于中位数，1=高于中位数

1. 这样做的目的是将 “预测具体房价” 转化为 “判断房价高低”。

分层抽样（Stratified Sampling）

• 使用StratifiedShuffleSplit确保训练集和测试集中高低房价的比例与原始数据一致：

sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_idx, test_idx in sss.split(X_scaled, y_category):
    X_train, X_test = X_scaled[train_idx], X_scaled[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]  # 注意：这里仍然使用原始的连续房价作为目标_