2
$\begingroup$

I'm trying to do a stratified split for a skewed dataset with target variable 'b'. The target variable is a bit value (either 0 or 1). Here's an example:

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

This code returns the following:

(93105, 2) (38, 2)

My problem is that the returned train/test arrays do not meet the set split ratio of 20%.

My setup:

  • Python 3.7.0 (32bit)
  • Sklearn 0.20.3
  • Pandas 0.23.4

I discovered that the problem is resulting from an integer overflow in the underlying split function.

How can I resolve this issue and is this a known bug? I couldn't find anything helpful.

$\endgroup$
3
  • 2
    $\begingroup$ The problem doesn't seem to apply to Python 3.7.x (64bit). $\endgroup$
    – tk78
    Apr 30, 2019 at 12:45
  • 1
    $\begingroup$ I'm glad to see that you(?) posted this as an issue on skl's github: github.com/scikit-learn/scikit-learn/issues/13794 $\endgroup$
    – Ben Reiniger
    Mar 25, 2020 at 21:46
  • $\begingroup$ @BenReiniger Yes, I did. $\endgroup$
    – tk78
    Apr 20, 2020 at 10:26

1 Answer 1

1
$\begingroup$

How about this?

# Split the data between the Training Data and Test Data
xTrain , xTest , yTrain , yTest = train_test_split(X , y , 
                                                  test_size = 0.30 , 
                                                  random_state = 0, 
                                          ----->  stratify = y)
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.