machine-learning - Scikit learn 中是否有库或工具用于将子类数据集拆分为文件夹中的训练/测试集

Question

我有一个不平衡8类数据集，每个类中有4个子类，并且我想将它分成测试/训练集的文件夹，数据集是这样的:

Dataset
   |------ Class A
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class B
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class C
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class D
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class E
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class F
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class G
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4
   |
   |------ Class K
   |  |---------Sub-Class 1
   |  |---------Sub-Class 2
   |  |---------Sub-Class 3
   |  |---------Sub-Class 4

因此，我正在寻找一种高效且专业的工具，可以帮助我将数据集拆分到训练/测试集的文件夹中。

Answer 1

我建议使用 train_test_split sklearn 库中的函数。它有参数 - 分层。对于此参数，您可以指定数据集中的几列进行分层。

例如:

X_train, X_test, y_train, y_test = train_test_split(Dataset, 
                                                    Dataset['Class'],
                                                    test_size=0.1, 
                                                    stratify=Dataset[['Sub-Class']])