如何拆分数据集以训练/测试某些行依赖的位置?

时间:2017-08-31 11:41:30

标签: python pandas train-test-split

我有一个主题数据集,每个主题在我的pandas数据框中有很多行(每个测量都是一行,一个主题可以测量几次)。我想将我的数据分成训练和测试集,但我不能随意拆分,因为所有受试者的测量都是相关的(不能将相同的主题放在火车和测试中)。你会如何重新获得这个?我有一个pandas数据帧,每个主题都有不同的测量数量。

编辑:我的数据包括每行的主题编号,我希望尽可能接近0.8 / 0.2。

1 个答案:

答案 0 :(得分:2)

考虑使用列df的数据框user_id来识别用户。

df = pd.DataFrame(
    np.random.randint(5, size=(100, 4)), columns=['user_id'] + list('ABC')
)

您想要识别唯一身份用户并随机选择一些用户。然后拆分数据框,以便将所有测试用户放在一个用户中,并在另一个用户中培训用户。

unique_users = df['user_id'].unique()
train_users, test_users = np.split(
    np.random.permutation(unique_users), [int(.8 * len(unique_users))]
)

df_train = df[df['user_id'].isin(train_users)]
df_test = df[df['user_id'].isin(test_users)]

这应该粗略地将您的数据拆分为80/20。

但是,如果您希望尽可能保持平衡,则必须逐步添加用户。

unique_users = df['user_id'].unique()
target_n = int(.8 * len(df))
shuffled_users = np.random.permutation(unique_users)

user_count = df['user_id'].value_counts()

mapping = user_count.reindex(shuffled_users).cumsum() <= target_n
mask = df['user_id'].map(mapping)

df_train = df[mask]
df_test = df[~mask]