如何使用pandas按标签分层以指定比例拆分训练集与测试集

你可以按以下纯pandas代码实现按标签比例分层拆分数据集，拆分比例设置为独立变量可随时调整：

import pandas as pd

# 独立配置项：训练集占比，后续调整拆分比例直接修改该值即可
TRAIN_RATIO = 0.8
# 随机种子，固定后每次拆分结果一致，无需复现可删除该参数
RANDOM_SEED = 42

# 你的原始数据构造逻辑
df = pd.DataFrame({'sentence': {0: 'FU bro',
  1: 'Well thats kinda cool',
  2: 'Haha thats so funny',
  3: 'cant u make somethin else mtfk',
  4: 'what a shame'},
 'label': {0: 'spam', 1: 'not spam', 2: 'not spam', 3: 'spam', 4: 'spam'}})

# 按标签拆分两类数据集
spam_df = df[df['label'] == 'spam']
not_spam_df = df[df['label'] == 'not spam']

# 分别对两类数据按比例抽取训练集，剩余数据作为测试集
spam_train = spam_df.sample(frac=TRAIN_RATIO, random_state=RANDOM_SEED)
spam_test = spam_df.drop(spam_train.index)

not_spam_train = not_spam_df.sample(frac=TRAIN_RATIO, random_state=RANDOM_SEED)
not_spam_test = not_spam_df.drop(not_spam_train.index)

# 合并得到最终训练集、测试集，额外做一次全量打乱避免标签聚集
train_df = pd.concat([spam_train, not_spam_train]).sample(frac=1, random_state=RANDOM_SEED).reset_index(drop=True)
test_df = pd.concat([spam_test, not_spam_test]).sample(frac=1, random_state=RANDOM_SEED).reset_index(drop=True)

# 输出验证拆分结果是否符合要求
print(f"训练集总条数：{len(train_df)}")
print(f"训练集spam条数：{len(train_df[train_df['label'] == 'spam'])}")
print(f"训练集not spam条数：{len(train_df[train_df['label'] == 'not spam'])}")
print(f"测试集总条数：{len(test_df)}")
print(f"测试集spam条数：{len(test_df[test_df['label'] == 'spam'])}")
print(f"测试集not spam条数：{len(test_df[test_df['label'] == 'not spam'])}")

补充可选方案 #

如果允许引入sklearn依赖，可以用更简洁的分层拆分接口实现完全一致的效果：

import pandas as pd
from sklearn.model_selection import train_test_split

TRAIN_RATIO = 0.8
RANDOM_SEED = 42

# 原始数据构造
df = pd.DataFrame({'sentence': {0: 'FU bro',
  1: 'Well thats kinda cool',
  2: 'Haha thats so funny',
  3: 'cant u make somethin else mtfk',
  4: 'what a shame'},
 'label': {0: 'spam', 1: 'not spam', 2: 'not spam', 3: 'spam', 4: 'spam'}})

# 直接调用分层拆分接口，stratify参数指定按label列保持比例拆分
train_df, test_df = train_test_split(
    df,
    train_size=TRAIN_RATIO,
    stratify=df['label'],
    random_state=RANDOM_SEED
)

关键说明 #

• 拆分比例TRAIN_RATIO为独立配置变量，后续需要调整拆分比例时仅需修改该值，无需改动其他逻辑
• 两种方案均采用分层抽样逻辑，保证训练集、测试集的标签占比和原数据集完全一致，针对你提供的100条数据（50条spam、50条not spam），TRAIN_RATIO=0.8时会自动得到80条训练集（40spam+40not spam）、20条测试集（10spam+10not spam），完全符合要求

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