如何按原数据集Y值比例拆分训练集与验证集？

我的数据集结构如下：

Title  Source Y
aaaaa  a      1
bbbbb  a      0
ccccc  b      0
ddddd  c      0
eeeee  c      0
fffff  a      0
ggggg  b      0
hhhhh  c      1
iiiii  a      0
jjjjj  a      0
....
....
....

其中Y为目标值，原数据集中Y=1占20%，Y=0占80%。

我当前使用以下代码拆分数据集（train_val_split=0.4）：

def split_dataset(self, dataset: Dataset | DatasetDict) -> Dataset | DatasetDict:
        if self.train_val_split is not None:
            split = dataset["train"].train_test_split(self.train_val_split)
            dataset["train"] = split["train"]
            dataset["validation"] = split["test"]
        dataset = self._select_samples(dataset)
        return dataset

拆分后训练集与验证集的Y值比例失衡：训练集含2个Y=1样本，验证集无Y=1样本。我需要拆分后训练集和验证集的Y=1占比均约为20%，预期拆分结果示例如下：

训练集：

Title  Source Y
ccccc  b      0
ddddd  c      0
eeeee  c      0
fffff  a      0
ggggg  b      0
hhhhh  c      1

验证集：

Title  Source Y
aaaaa  a      1
bbbbb  a      0
iiiii  a      0
jjjjj  a      0

请问是否有方法实现这种按原比例的数据集拆分？

可以通过分层抽样实现，核心是让拆分过程参考目标列Y的类别分布，确保训练集和验证集都保留原数据集的类别比例。针对你使用的Hugging Face Dataset类，只需修改train_test_split的调用参数即可：

修改后的代码：

def split_dataset(self, dataset: Dataset | DatasetDict) -> Dataset | DatasetDict:
        if self.train_val_split is not None:
            # 按Y列分层抽样，保证类别比例一致
            split = dataset["train"].train_test_split(
                test_size=self.train_val_split,
                stratify_by_column="Y"
            )
            dataset["train"] = split["train"]
            dataset["validation"] = split["test"]
        dataset = self._select_samples(dataset)
        return dataset

核心要点 #

• stratify_by_column="Y"：该参数指定以Y列为分层依据，拆分时会按原数据中Y=0和Y=1的比例，将样本分配到训练集和验证集，避免出现类别失衡。
• 参数test_size对应你原本的train_val_split，代表验证集占原训练集的比例，和你的需求完全匹配。

这种方法同样适用于多分类场景，只需将stratify_by_column指定为对应的目标列即可。

内容的提问来源于stack exchange，提问作者code12345