Hugging Face Datasets train_test_split 怎样按标签分层拆分训练集与测试集

如果一个带标签的数据集直接调用 train_test_split(test_size=0.2)，默认只是随机抽取测试集；当类别分布不均衡时，测试集可能偶然偏向某个标签。要让切分尽量保持类别比例，应把标签列名传给 stratify_by_column，再用固定的 seed 让同一份输入可以复现。

官方地址：https://huggingface.co/docs/datasets/

下面用人工构造的短文本和整数标签演示完整流程。示例只解释 API 的参数关系，不把图示或代码片段冒充已经在本机执行的输出。

先确认标签列和切分目标

开始前先回答三个问题：哪一列代表类别、测试集需要占多少比例、是否需要每次得到相同的行划分。stratify_by_column 接收的是列名，不是列中某一个标签值；test_size 可以使用比例，也可以使用固定行数。

from datasets import Dataset

# 使用人工样本演示二分类；真实项目替换为自己的数据加载过程。
rows = {
    "text": [
        "样本一：设备运行正常。",
        "样本二：发现轻微抖动。",
        "样本三：设备运行正常。",
        "样本四：发现轻微抖动。",
        "样本五：设备运行正常。",
        "样本六：发现轻微抖动。",
        "样本七：设备运行正常。",
        "样本八：设备运行正常。",
        "样本九：发现轻微抖动。",
        "样本十：设备运行正常。",
    ],
    # 0 和 1 是演示标签；生产数据应先确认标签编码含义。
    "label": [0, 1, 0, 1, 0, 1, 0, 0, 1, 0],
}
dataset = Dataset.from_dict(rows)

# 先确认列名和样本量，避免把 text 写成标签列。
print(dataset.column_names)
print(dataset.num_rows)

这里的标签列叫 label，而且每个类别都有多条样本。若标签列不存在、列名拼写错误，或者某个类别只有一条记录，分层切分就可能失败；这不是把 test_size 调小就一定能修好的问题。

用 stratify_by_column 执行分层切分

把 stratify_by_column="label" 传入后，Datasets 会使用该列参与分层抽样，返回一个包含 train 和 test 的 DatasetDict。它仍然可能受到总行数、标签类别数量和取整方式影响，所以“保持比例”应理解为在可分配条件下尽量保持，而不是承诺每个 split 的小数比例完全相等。

# test_size=0.2 表示把约两成样本放入测试集。
# stratify_by_column 指向列名，让切分参考每条记录的 label。
splits = dataset.train_test_split(
    test_size=0.2,
    stratify_by_column="label",
    seed=42,
)

# 返回对象按名称访问两个 Dataset split。
train_ds = splits["train"]
test_ds = splits["test"]
print(train_ds.num_rows, test_ds.num_rows)
print(splits.keys())

不要把返回值当成两个普通列表，也不要继续对原始 dataset 误读为“已经切分”。后续训练使用 splits["train"]，评估使用 splits["test"]，这样数据边界会留在返回的 DatasetDict 中。

用 seed 固定可复现边界

seed 控制随机过程。相同的数据内容、相同的参数和相同的种子，才有机会得到同一组行划分；如果重新加载的数据顺序已变、样本有增删，固定种子也不能把不同输入变成同一结果。

def split_once(ds):
    # 将所有影响切分的参数集中，避免训练脚本和评估脚本各写一套。
    return ds.train_test_split(
        test_size=0.2,
        shuffle=True,
        stratify_by_column="label",
        seed=2026,
    )

first = split_once(dataset)
second = split_once(dataset)

# 用稳定字段比较前几行；这里只检查切分逻辑，不比较对象地址。
assert first["test"]["text"] == second["test"]["text"]
assert first["test"]["label"] == second["test"]["label"]

生产流水线可以把种子写进配置或实验记录，并同时记录原始数据版本、列名和 test_size。只保存一个整数种子而不保存输入数据快照，仍不足以完整复盘一次实验。

检查比例、稀有标签与失败原因

分层切分的排查顺序应该先看列，再看每类数量，最后看 split 结果。不要只看总行数相加是否等于原数据，因为那只能证明没有明显丢行，不能证明类别分布符合预期。

from collections import Counter

def label_counts(ds):
    # Counter 让训练集和测试集的标签构成可以直接对照。
    return Counter(ds["label"])

print("all:", label_counts(dataset))
print("train:", label_counts(splits["train"]))
print("test:", label_counts(splits["test"]))

# 先确认两部分合计仍等于原始样本数。
assert splits["train"].num_rows + splits["test"].num_rows == dataset.num_rows

把参数写成一份可交接配置

最终建议把切分决策和数据快照一起记录：标签列叫什么、测试集比例是多少、是否打乱、使用了哪个种子、每个 split 的行数和标签计数是多少。这样换机器或重跑训练时，问题会落在明确的输入差异上，而不是一句“随机切分不稳定”。

简要记忆：默认随机切分解决的是“分成两份”，stratify_by_column 解决的是“尽量保持标签构成”，seed 解决的是“相同输入下可复现”。三者职责不同，不能用固定种子替代分层，也不能用分层参数替代数据版本记录。

只有一个标签类别时还能做分层切分吗？

技术上可以切分行，但分层没有提供类别平衡价值。若任务本来就是单类别数据，应明确记录这一事实，而不是把 stratify_by_column 当作通用稳定性开关。

test_size 能保证测试集严格占 20% 吗？

它表达目标比例或固定数量，实际行数还会受总行数、整数取整和标签可分配条件影响。小数据集更应该检查最终行数和类别计数。

固定 seed 后为什么换数据文件仍然不同？

种子只固定随机过程，不固定输入内容。样本增删、顺序改变、过滤条件改变，都会让相同随机过程作用在不同序列上。