python - 带有分层的 Test_train_split

Question

我正在尝试按数据帧(~188k 行)拆分为训练样本和测试样本。列 ('FLAG') 是我的目标变量，包含值 0 或 1。

由于只有大约 1300 个值为 1 的“FLAG”，我想进行分层拆分以确保两个样本中都有代表性数量的 1 值。

我尝试使用 sklearn 的 train_test_split 函数进行拆分:

train, test = train_test_split(df, test_size=0.2, stratify=df["FLAG"])

我的问题是，生成的训练和测试样本分别有 177942 行和 52 行。我本以为会有 150400 和 37600 行。

我阅读文档 (sklearn.model_selection.train_test_split) 的理解是，我必须提供我的数据框、test_size 和包含目标类的列(即在我的例子中为“FLAG”)。

即使是一个通用的例子:

df = pd.DataFrame(data={'a': np.random.rand(100000), 'b': np.random.rand(100000), 'c': 0})
df.loc[np.random.randint(0, 100000, 1000), 'c'] = 1
tr, ts = train_test_split(df, test_size=.2, stratify=df['c'])
print(tr.shape, ts.shape)

返回:(93105, 3) (38, 3)

我的进口 list :

import cx_Oracle
import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np

我的python版本:3.7.0 sklearn版本:0.20.3 Pandas 版本:0.23.4

Answer 1

我的调查表明该问题是由整数溢出引起的。该问题仅发生在 Python 3.7.x 32 位上。 64 位版本工作正常。

最后我切换到 64 位 Python 来解决这个问题(我之前不得不使用 32 位版本，因为不相关的 Oracle 包依赖)。