如何进行训练测试拆分，以便 Python 中每个类都有足够的训练和测试数据？答案

【问题标题】：How to do train test split such that there are enough training and testing data from each class in Python?如何进行训练测试拆分，以便 Python 中每个类都有足够的训练和测试数据？
【发布时间】：2019-10-11 18:19:50
【问题描述】：

我有一个数据集，它有 5 个类别，分布如下：

从分布中可以明显看出，1 类的样本非常少。

如何对这些数据进行训练-测试拆分，以便 Python 中每个类别都有足够的训练和测试数据？

【问题讨论】：

您可能会获得每个类别的 70% .. 像 df.groupby('category').apply(train_test_split, args=0.7, ....) 这样您填写正确的参数而不是 args 等
通常，当您执行df.sample 时，它会随机抽取您的训练集。从理论上讲，这将为您提供相同的课程分布。也就是说，你总是可以加倍努力，按照@rafaelc 的建议去做。

标签： python pandas scikit-learn

【解决方案1】：

train_test_split 函数默认在拆分前对数据集进行洗牌，除非您将 shuffle 参数值提供为 False。如果shuffle 为真，我认为它可以确保您训练数据集的部分将具有来自所有类别的值。此外，如果您希望 train_test_split 的结果具有确定性，则可以使用 random_state 参数。请参阅documentation 了解更多信息。希望对您有所帮助。

【讨论】：

我认为您提供的答案不能满足 OP 的要求，因为train_test_split(...., test_size=0.2, shuffle=False) 不能确保从每个班级中平等抽取样本。您可能会找到通过正确切分数据来手动拆分数据的方法，然后将其提供给您的 NN 或在其上使用 ML。
由于每个类别的训练样本大小可能会有所不同，因此不可能在每种情况下都从每个类别中平均分割训练数据集。我提到它将确保从每个类别中至少抽取一些样本，其中我的意思不是平等。谢谢。
@RubelHassan，因为您接受 由于每个类别的训练样本的大小可能会有所不同，因此不可能在每种情况下都从每个类别中平均分割训练数据集，那么它如何回答OP的问题？请注意，这个问题与一个名为不平衡数据的概念有关。因此，建议您修改建议的答案或将其删除，以免被否决。
@mnm OP 的问题如何对这些数据进行训练-测试拆分，以便在 Python 中每个类别都有 _enough_ 训练和测试数据？，你能看到他在哪里要求平等。

【解决方案2】：

将train_test_split 中的stratify 参数设置为您的目标列。

stratify 将确保每个班级平均分配。 Doc

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, stratify=y)

【讨论】：