I'm working on a dataset that isn't split into test and train set by default and I'm a bit concerned about the imbalance between the 'label' distributions between them and how they might affect the trained model's performance. Let me note that I use deep neural networks and the prediction type is regression.

By sequentially splitting the samples into test/train (20/80) I get the following distributions respectively.

test labels train labels

I'm worried, since model performance is not improving by tuning hyperparamaters, if I'm generally allowed to try random seeds for shuffling my dataset till test/train label distributions look alike. Are there any references/best practices for that? I'm not concerned about compromising the ability to generalize/overfitting since I don't base my splitting into the form of the input data, rather the predicted outcome.


1 Answer 1


Why don't you simply stratify your train test split? Sklearn provides a built-in function/option (stratify=) to do this:

from sklearn.model_selection import train_test_split
Xtrain, Xtest, ytrain, ytest = train_test_split(X, y, stratify=y, test_size=0.2, random_state=12)
  • $\begingroup$ That's very interesting, I'm relieved to know this a documented practice. I wonder if it also applies to continuous variable outcomes for regression problems which fit my case. $\endgroup$
    – civy
    Feb 9 at 17:26
  • $\begingroup$ Ah, I see. Usually when you draw a random sample from some data, you should be able to "reproduce" the original distribution as long as your sample is not too small. So with continuous data, stratification is seldom used. However, you could assign "bins" to y-values and stratify according to these bins to "reproduce" the original distribution. kaggle.com/tolgadincer/continuous-target-stratification $\endgroup$
    – Peter
    Feb 9 at 18:36

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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