将训练/测试拆分为精确的行数

Question

有一个大约 6,000,000 行的数据集。我使用以下代码将数据集分为训练集和测试集：

from sklearn.model_selection import train_test_split
train, test = train_test_split(df, test_size=0.3, random_state=42)

我想要的是将测试集精确地拆分为一定数量的行，例如测试集 50,000 行。

有什么方法可以得到这个？

Answer 1

根据帮助页面：

test_size float 或 int，default=None 如果是 float，则应介于 0.0 和 1.0 之间，并表示要包含在测试拆分中的数据集的比例。如果是int，则表示测试样本的绝对数量。如果无，则该值设置为列车大小的补数。如果train_size也为None，它将被设置为0.25。

您可以为 test_size 指定一个整数：

train, test = train_test_split(df, test_size=50000, random_state=42)