### 两者的区别 | 对比项 | `train_test_split` | `StratifiedShuffleSplit` | |-----|------|------| | 分割方式 | 纯随机抽样，类似掷骰子 | **分层抽样**，按类别比例抽取 | | 类别分布 | 小类别占比可能波动较大（如 2% 的类可能只剩 1.2% 或 2.8%） | 每个类别在训练集/测试集中的比例与原始数据**严格一致** | | 调用方式 | 一次调用直接返回切分后的数据 `X_train, X_test, ...` | 返回的是**索引**，需自己用 `.iloc` 取数据 | | 多次分割 | 每调用一次切一次 | 通过 `n_splits` 可一次性生成多组切分（类似交叉验证） | | 适用场景 | 数据量大、类别均衡 | **类别不平衡**（如意图识别中"退款咨询"占 25%，"技术支持"仅占 2%）、数据量较少 | **核心差异说明**： - `train_test_split` 是"盲切"：1000 条数据按 70/15/15 切完，占比 2% 的"技术支持"（约 20 条）在测试集（150 条）里理论上只有 3 条，随机波动下可能是 1 条或 5 条，占比偏差可达 ±50% 以上，导致评估指标不稳定。 - `StratifiedShuffleSplit` 是"按比例分层切"：先按每个类别的比例算好测试集里该放几条，再在各类别内部随机打乱抽取，保证 2% 的类在测试集中也精确占 2%（3 条），训练/验证/测试三个集合的类别分布与原始分布几乎无差异。 - 事实上 `train_test_split` 也有 `stratify=y` 参数可以达到同样效果，两者底层逻辑一致；`StratifiedShuffleSplit` 的优势是支持 `n_splits` 多组切分和更灵活的 API 形式（返回索引）。

### StratifiedShuffleSplit 参数详解 ```python class sklearn.model_selection.StratifiedShuffleSplit( n_splits=10, # 分割次数 *, # 以下为仅关键字参数 test_size=None, # 测试集比例或数量 train_size=None, # 训练集比例或数量 random_state=None # 随机种子 ) ``` | 参数 | 类型 | 默认值 | 说明 | |-----|------|------|------| | `n_splits` | int | 10 | **重新分割并迭代的次数**。每次 `split()` 调用会生成一组不同的训练/测试索引，共生成 `n_splits` 组。设为 1 表示只切一次（普通的一次性分割）；设为 5 则类似 5 折分层交叉验证（但每组切分相互独立，不保证互斥） | | `test_size` | float 或 int | None | 测试集大小。float（0~1）表示**比例**，如 `0.3` 表示 30% 作为测试集；int 表示**绝对条数**，如 `150` 表示取 150 条。若 `test_size` 和 `train_size` 都为 None，默认 `test_size=0.1` | | `train_size` | float 或 int | None | 训练集大小，含义同上。两者只需设置一个，另一个自动取补集；也可以同时设置（比例之和为 1），中间剩余部分既不进训练集也不进测试集（可用于留出中间集） | | `random_state` | int 或 RandomState | None | 随机种子。固定后每次运行结果完全一致（可复现）；不设则每次切分结果不同 | **`split(X, y)` 方法**（不是构造参数，但必须掌握）： - `y` 必须提供，且是**离散类别标签**——分层就是按 `y` 的类别比例来分的，这是与 `train_test_split` 最不同的地方（后者的 `y` 可有可无）。 - `split()` 返回一个**生成器**，每次 `next()`（或 for 循环）产出一组 `(train_indices, test_indices)`，是 numpy 索引数组而非数据本身，所以要用 `.iloc` 取回数据。 - 每个类别的样本数必须不少于 2（否则无法分层），且测试集/训练集大小不能小于类别数（每个类至少要在每边分到 1 条）。 **本例中的参数用法**： ```python # 第一次：70% 训练 / 30% 暂存（后续再切验证+测试） sss = StratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=42) train_idx, temp_idx = next(sss.split(X, y)) # 取生成器的第 1 组索引 # 第二次：暂存集对半切，15% 验证 / 15% 测试 sss2 = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42) val_idx, test_idx = next(sss2.split(X_temp_strat, y_temp_strat)) ``` - `n_splits=1`：只需要一组切分，用 `next()` 取出即可； - `random_state=42`：保证每次运行代码切出的都是同一批数据； - `next(...)`：因为 `split()` 返回生成器，`next()` 拿到第 1 组（也是唯一一组）`(train索引, test索引)`。

### 数据分割实战代码 ```python # ✂️ 数据分割实战代码 from sklearn.model_selection import train_test_split, StratifiedShuffleSplit import pandas as pd import numpy as np import matplotlib.pyplot as plt # 使用前面清理好的数据，这里创建一个更大的模拟数据集 np.random.seed(42) # 模拟智能客服数据集 intents = ['退款咨询', '物流查询', '优惠咨询', '售后投诉', '联系方式', '换货咨询', '使用咨询', '产品咨询', '订单查询', '技术支持'] # 模拟不平衡的数据分布（符合实际情况） intent_weights = [0.25, 0.20, 0.15, 0.12, 0.08, 0.06, 0.05, 0.04, 0.03, 0.02] # 生成模拟数据 data_size = 1000 texts = [] labels = [] for i in range(data_size): intent = np.random.choice(intents, p=intent_weights) # 简单模拟文本 text = f"这是一条关于{intent}的用户问询{i}" texts.append(text) labels.append(intent) df = pd.DataFrame({ 'text': texts, 'intent': labels }) print("📊 === 原始数据概览 ===") print(f"总数据量: {len(df)} 条") print("\n各意图分布:") intent_counts = df['intent'].value_counts() for intent, count in intent_counts.items(): percentage = count / len(df) * 100 print(f" {intent}: {count:3d}条 ({percentage:4.1f}%)") # 可视化原始数据分布 plt.figure(figsize=(12, 4)) plt.subplot(1, 2, 1) intent_counts.plot(kind='bar', color='skyblue', alpha=0.7) plt.title('📊 原始数据分布', fontsize=12, weight='bold') plt.xlabel('意图类别') plt.ylabel('数量') plt.xticks(rotation=45) print("\n✂️ === 方法1: 简单随机分割 ===") # 方法1: 简单随机分割 X = df['text'] y = df['intent'] X_train, X_temp, y_train, y_temp = train_test_split( X, y, test_size=0.3, 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) print(f"训练集: {len(X_train)} 条 ({len(X_train)/len(df)*100:.1f}%)") print(f"验证集: {len(X_val)} 条 ({len(X_val)/len(df)*100:.1f}%)") print(f"测试集: {len(X_test)} 条 ({len(X_test)/len(df)*100:.1f}%)") print("\n✂️ === 方法2: 分层分割（推荐）===") # 方法2: 分层分割，保证各类别比例一致 sss = StratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=42) train_idx, temp_idx = next(sss.split(X, y)) X_train_strat = X.iloc[train_idx] y_train_strat = y.iloc[train_idx] X_temp_strat = X.iloc[temp_idx] y_temp_strat = y.iloc[temp_idx] # 再次分割验证集和测试集 sss2 = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42) val_idx, test_idx = next(sss2.split(X_temp_strat, y_temp_strat)) X_val_strat = X_temp_strat.iloc[val_idx] y_val_strat = y_temp_strat.iloc[val_idx] X_test_strat = X_temp_strat.iloc[test_idx] y_test_strat = y_test.iloc[test_idx] # 可视化分割结果比较 plt.subplot(1, 2, 2) comparison_data = { '原始': intent_counts.head(5), '训练集': y_train_strat.value_counts().head(5), '测试集': y_test_strat.value_counts().head(5) } x = np.arange(5) width = 0.25 for i, (name, data) in enumerate(comparison_data.items()): plt.bar(x + i*width, data.values, width, label=name, alpha=0.8) plt.title('📊 分层分割效果对比', fontsize=12, weight='bold') plt.xlabel('意图类别（前5个）') plt.ylabel('数量') plt.xticks(x + width, data.index, rotation=45) plt.legend() plt.tight_layout() plt.show() ``` > 注意：示例中 `y_test_strat = y_test.iloc[test_idx]` 应为 `y_temp_strat.iloc[test_idx]`，`y_test` 是随机分割的结果，二者混用会报索引对齐错误。

### 分割效果与总结 **分割质量验证**：用"与原始分布的平均绝对差异"量化分割质量（越小越好）： ```python def calculate_distribution_difference(original, subset): """计算分布差异""" orig_dist = original.value_counts(normalize=True).sort_index() subset_dist = subset.value_counts(normalize=True).sort_index() # 确保所有类别都存在 all_classes = orig_dist.index.union(subset_dist.index) orig_dist = orig_dist.reindex(all_classes, fill_value=0) subset_dist = subset_dist.reindex(all_classes, fill_value=0) # 计算平均绝对差异 diff = np.mean(np.abs(orig_dist - subset_dist)) return diff print("与原始分布的差异度（越小越好）:") print(f"随机分割 - 训练集: {calculate_distribution_difference(y, y_train):.4f}") print(f"随机分割 - 测试集: {calculate_distribution_difference(y, y_test):.4f}") print(f"分层分割 - 训练集: {calculate_distribution_difference(y, y_train_strat):.4f}") print(f"分层分割 - 测试集: {calculate_distribution_difference(y, y_test_strat):.4f}") ``` **典型结果对比**： | 分割方式 | 训练集差异度 | 测试集差异度 | |-----|---------|---------| | 随机分割 | ~0.005 - 0.01（个别小类偏差明显） | ~0.01 - 0.02（小类波动更大） | | 分层分割 | ~0.000 - 0.002（几乎无差异） | ~0.000 - 0.002（几乎无差异） | **总结**： - 🎯 **推荐方案**：分层分割 - 📊 **分割比例**：70% 训练 + 15% 验证 + 15% 测试 - ✅ **优势**：保证了各类别在不同集合中的分布一致性，尤其保护占比小的类别（如 2% 的"技术支持"）不被随机波动"稀释"或"放大" - ✅ **结果**：模型训练更稳定，评估更可靠 - ✅ 训练集用于模型训练、验证集用于超参数调优、测试集用于最终评估，验证集和测试集都必须用分层方式切分