多标签图像分类加权训练/验证/测试集划分（PyTorch实现）

问题分析 #

你面对的是典型的多标签数据集分层拆分问题：单样本可关联多个疾病标签，常规单标签分层抽样无法同时兼顾所有类别的分布均衡，尤其当前数据集存在多个小样本类别（Fibrosis仅25例、Hernia仅1例），随机拆分很容易出现稀有类别仅分布在某一个子集的问题，直接影响模型验证和测试结果的可信度。

拆分方案选择 #

直接用**迭代分层法（Iterative Stratification）**即可，这是目前多标签数据集拆分的主流方案：算法每一轮优先分配带有「当前分布偏差最大的稀有标签」的样本，将其放到该标签占比最低的子集，迭代直到所有样本分配完成，能最大程度让三个子集的每一类疾病占比和原数据集保持一致。

注：Hernia类别仅1个样本，不存在任何拆分方法能让三个子集同时包含该类样本，若需要评估模型对该类的识别能力，需先补充该类样本。

PyTorch 实现代码 #

依赖安装 #

pip install pandas numpy torch scikit-multilearn pillow
scikit-multilearn内置了迭代分层拆分的成熟实现，不需要自己手写算法逻辑。

完整实现步骤 #

import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from skmultilearn.model_selection import iterative_train_test_split
from PIL import Image

# ----------------------
# 1. 配置全局参数
# ----------------------
# 所有疾病类别，和你给出的统计结果一一对应
DISEASE_LABELS = [
    'Atelectasis', 'Infiltration', 'No Finding', 'Emphysema', 'Pneumonia',
    'Pleural_Thickening', 'Pneumothorax', 'Mass', 'Nodule', 'Consolidation',
    'Edema', 'Cardiomegaly', 'Effusion', 'Fibrosis', 'Hernia'
]
LABEL2IDX = {lab: i for i, lab in enumerate(DISEASE_LABELS)}
N_CLASSES = len(DISEASE_LABELS)
BATCH_SIZE = 32

# 加载你自己的数据集DataFrame，替换成你的实际路径
# df = pd.read_csv("your_dataset_metadata.csv")

# ----------------------
# 2. 多标签分层拆分
# ----------------------
# 构造多热标签矩阵，形状为(样本总数, 类别数)
y = np.zeros((len(df), N_CLASSES), dtype=np.int8)
for i, label_str in enumerate(df['Finding_Labels'].values):
    # 按你数据集实际的标签分隔符切分，截图中为|分隔
    for lab in label_str.split('|'):
        if lab in LABEL2IDX:
            y[i, LABEL2IDX[lab]] = 1

# 特征用样本索引即可，后续靠索引取对应的数据
X = np.arange(len(df)).reshape(-1, 1)

# 第一次拆分：80%训练集，剩余20%为待拆分的临时集
X_train, y_train, X_temp, y_temp = iterative_train_test_split(X, y, test_size=0.2)
# 第二次拆分：临时集对半分，得到10%验证集、10%测试集
X_val, y_val, X_test, y_test = iterative_train_test_split(X_temp, y_temp, test_size=0.5)

# 提取各子集对应的DataFrame行索引
train_idx = X_train.flatten()
val_idx = X_val.flatten()
test_idx = X_test.flatten()

# 可选：校验各子集标签分布
def check_dist(label_matrix, label_names):
    total = len(label_matrix)
    return {name: round(label_matrix[:,i].sum()/total*100, 2) for i,name in enumerate(label_names)}
print("原集标签占比%：", check_dist(y, DISEASE_LABELS))
print("训练集标签占比%：", check_dist(y_train, DISEASE_LABELS))
print("验证集标签占比%：", check_dist(y_val, DISEASE_LABELS))
print("测试集标签占比%：", check_dist(y_test, DISEASE_LABELS))

# ----------------------
# 3. 对接PyTorch数据加载逻辑
# ----------------------
class ChestXRayDataset(Dataset):
    def __init__(self, sub_df, transform=None):
        self.df = sub_df.reset_index(drop=True)
        self.transform = transform
        # 预加载子集标签，避免训练时重复计算
        self.labels = np.zeros((len(self.df), N_CLASSES), dtype=np.float32)
        for i, label_str in enumerate(self.df['Finding_Labels'].values):
            for lab in label_str.split('|'):
                if lab in LABEL2IDX:
                    self.labels[i, LABEL2IDX[lab]] = 1.0

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 替换成你实际的图像读取逻辑，Image_Index对应当前df里的图像路径列
        img = Image.open(self.df.iloc[idx]['Image_Index']).convert('RGB')
        if self.transform:
            img = self.transform(img)
        label = torch.tensor(self.labels[idx])
        return img, label

# 生成各子集数据集，可自行添加训练/验证对应的图像增强
# train_transform = ... # 训练集增强：随机裁剪、翻转、归一化等
# val_transform = ... # 验证/测试集变换：仅 resize、归一化
train_dataset = ChestXRayDataset(df.iloc[train_idx])
val_dataset = ChestXRayDataset(df.iloc[val_idx])
test_dataset = ChestXRayDataset(df.iloc[test_idx])

# 生成DataLoader
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

补充说明 #

• 正常拆分后，样本量大于100的类别，三个子集的占比和原数据集偏差可以控制在0.5%以内，小样本类别偏差也不会超过2%。
• 若不想引入scikit-multilearn依赖，可参考sechidis等人2011年提出的迭代分层原理论文自行实现，但直接用库的成熟实现稳定性更高，拆分速度也更快。
• 拆分完成后必须做标签分布校验，避免拆分逻辑bug导致分布偏差过大。

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