python - 如何按比例将数据分为训练集和测试集？

Question

如何编写一个采用数据集（我正在使用虹膜）并创建一个测试和数据集的函数，并按比例表示每个类别（标签）？

def train_test_split(data, labels, n, test_proportion)

其中n是类别数

输出的test_data应该包含数据的test_proportion％，其中test_labels包含与test_data中的数据相对应的正确标签；然后train_data应该包含其余数据，train_labels包含train_data中特征向量的标签。

test_data和train_data在n个类别中都应具有相等的比例。例如，对于n = 3，则test_data和train_data都包含类别0的1/3观测值，类别1的1/3和类别2的1/3，即使test_data和train_data可能包含不同数量的条目（当test_proportion与0.5不同）。

Answer 1

我相信您要寻找的答案在于SKLearn的Train_Test_Split实现中：

https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html

也许更具体地说：

https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedShuffleSplit.html#sklearn.model_selection.StratifiedShuffleSplit