图像分类的加权训练、测试、验证拆分

Question

我正在尝试将上述 dataframe 拆分为训练（80%）、验证（10%）和测试（10%）； 但是，我想在每组中保持几乎相同数量的疾病。 Finding_Labels 列包含图像链接到的疾病列表。

如第一行所示，每张图像都可能与一种以上的疾病相关联——这有点问题。 因此，如何以 3 组疾病数量几乎相等的方式拆分它？

PyTorch 中的答案将不胜感激。

疾病名称及计数：

{'Atelectasis': 391,
         'Infiltration': 1181,
         'No Finding': 3479,
         'Emphysema': 116,
         'Pneumonia': 116,
         'Pleural_Thickening': 130,
         'Pneumothorax': 361,
         'Mass': 241,
         'Nodule': 190,
         'Consolidation': 299,
         'Edema': 124,
         'Cardiomegaly': 157,
         'Effusion': 506,
         'Fibrosis': 25,
         'Hernia': 1}

Answer 1

正如评论中提到的，我会使用sklearn来解决您的问题。

例如，假设您有以下特征和 label 数据集，其中一个属性的数量(y==1).sum()等于零属性的数量(y==0).sum() （所以你一对零的比率为1.0 ）：

from sklearn import datasets

X, y = datasets.make_blobs(n_samples=10000, centers=2, random_state=0)

print(np.bincount(y)[0]/np.bincount(y)[1]) # 1.0

然后，您可以将数据拆分为训练、验证和测试数据，同时保持相同的属性比率，如下所示：

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42, test_size=0.2)
X_val, X_test, y_val, y_test = train_test_split(X_test, y_test, stratify=y_test, random_state=42, test_size=0.5)

print(np.bincount(y_train)[0]/np.bincount(y_train)[1]) # 1.0
print(np.bincount(y_test)[0]/np.bincount(y_test)[1])   # 1.0
print(np.bincount(y_val)[0]/np.bincount(y_val)[1])     # 1.0
print(len(y_train), len(y_val), len(y_test))           # 8000, 1000, 1000