在机器学习中,模型需要经历两个核心阶段:
类比场景:学生通过 “练习题”(训练集)学习知识,再通过 “考试题”(测试集)检验真实水平。
train_test_split
函数的核心参数与逻辑X_train, X_test, y_train, y_test = train_test_split(
X_scaled, y, test_size=0.2, random_state=42
)
X_scaled
:特征矩阵(已标准化的面积、房龄等特征)。y
:目标变量(房价)。test_size=0.2
:测试集占总数据的比例(20%),也可设为整数(如 test_size=20
表示取 20 个样本)。random_state=42
:随机种子,确保每次划分结果一致(与 np.random.seed(42)
作用类似)。test_size
比例从原始数据中随机抽取样本作为测试集,剩余作为训练集。X
和 y
的样本顺序一一对应(如第 i 个特征向量对应第 i 个房价标签)。假设原始数据有 100 个样本(n_samples=100
):
X_train.shape=(80, 2)
, y_train.shape=(80,)
),用于模型学习。X_test.shape=(20, 2)
, y_test.shape=(20,)
),用于评估模型泛化能力。test_size
:平衡训练与测试的样本量test_size=0.2~0.3
(20%-30% 作为测试集)。test_size=0.1
甚至更低(因少量样本已足够评估)。test_size=1.0
,则所有数据都是测试集,无训练集;若 test_size=0
,则全是训练集。random_state
:确保可复现的 “随机” 划分random_state
:每次划分结果不同,导致模型评估指标波动。random_state=42
:多次运行代码,划分结果一致,便于对比不同模型效果。shuffle=True
(默认参数):打乱数据顺序train_test_split
默认为 shuffle=True
,即先打乱数据再划分;若数据已随机排列,可设 shuffle=False
。当目标变量是分类变量(如二分类 “是否违约”)时,普通随机划分可能导致训练 / 测试集的类别比例失衡(如测试集全是 “违约” 样本)。此时需用 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)}")
scaler.fit(X_train)
),最后用训练集的标准化参数转换测试集(scaler.transform(X_test)
)。通过合理划分训练集与测试集,你可以更准确地评估模型的实际能力,避免被 “过拟合” 的假象误导 —— 这是机器学习工程化中至关重要的一步!
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
数据准备与二分类转换
y
转换为二分类标签:threshold = np.median(y) # 使用中位数作为阈值
y_binary = (y > threshold).astype(int) # 0=低于中位数,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] # 注意:这里仍然使用原始的连续房价作为目标_