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, 2022 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, 2022 at 18:36

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.