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

Question

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

​

​

从发行版中可以明显看出，类1的样本非常少。

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

Answer 1

将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)

Answer 2

默认情况下，除非将shuffle参数值设置为False，否则train_test_split函数会在拆分之前打乱数据集。我认为，如果shuffle为真，它将确保数据集的训练部分具有所有类别的值。此外，如果您希望train_test_split的结果是确定性的，那么可以使用random_state参数。请参阅documentation了解更多信息。希望能有所帮助。