多输出图像分类下带类别样本阈值的多标签分层数据集拆分问题

核心逻辑 #

要实现同时对两个标签做分层拆分，同时支持自定义每个类别的样本量阈值，只需要两步改造原有单标签拆分逻辑：

• 先做阈值过滤：提前统计两个标签每个类别的总样本量，剔除低于阈值的类别样本，避免拆分后小类别样本完全缺失
• 生成联合分层标签：将两个目标标签拼接为单个新标签，以此作为分层依据，即可同时保证两个标签的分布在训练、验证集与原数据集一致

完整实现代码 #

import pandas as pd
from sklearn.model_selection import StratifiedShuffleSplit

# -------------------------- 自定义参数 --------------------------
# 两个需要分层的标签列名
tag_col1 = 'layout'
tag_col2 = 'replace_with_your_second_tag'
# 两个标签对应类别的最小样本阈值（总样本量低于该值的类别会被过滤）
tag1_min_count = 10
tag2_min_count = 8
# 验证集占比
val_size = 0.2
random_seed = 42
# ----------------------------------------------------------------

# 步骤1：过滤不符合样本量阈值的类别
tag1_valid = tdf[tag_col1].value_counts()[lambda x: x >= tag1_min_count].index
tag2_valid = tdf[tag_col2].value_counts()[lambda x: x >= tag2_min_count].index
filtered_df = tdf[
    tdf[tag_col1].isin(tag1_valid) & tdf[tag_col2].isin(tag2_valid)
].reset_index(drop=True)

# 步骤2：生成联合分层标签（分隔符可自定义，只要不与标签内容重复即可）
filtered_df['stratify_key'] = filtered_df[tag_col1].astype(str) + "||" + filtered_df[tag_col2].astype(str)

# 步骤3：执行分层拆分
split = StratifiedShuffleSplit(n_splits=1, test_size=val_size, random_state=random_seed)
for train_idx, val_idx in split.split(filtered_df, filtered_df['stratify_key']):
    train_df = filtered_df.loc[train_idx].drop(columns=['stratify_key'])
    val_df = filtered_df.loc[val_idx].drop(columns=['stratify_key'])

补充说明 #

• 如果不需要剔除低样本类别，仅要求拆分后训练/验证集的各标签类别满足样本量阈值，可以把n_splits参数调整为10~20，遍历所有生成的拆分结果，筛选出符合阈值要求的拆分对即可
• 该方案支持扩展到3个及以上标签的分层拆分，只需在生成stratify_key时拼接更多标签列即可
• 拆分完成后可分别统计两个标签在训练集、验证集的类别分布，和原数据集对比，分层误差通常低于1%，符合业务要求

内容的提问来源于stack exchange，提问作者Tùng Dương Quang