I faced a problem while I using sklearn.train_test_split(). Here is the code I use.

xtrain, xtest, ytrain, ytest = train_test_split(X_source, Y_source, test_size=0.3)

The shape of X_source is (2427,features_size), Y_source is (2427,1). And there are 65 different classes of labels in Y_source, what I mean is that Y_source is a matrix length of 2427 and the value from 1 to 65.

My problem is that, while I use train_test_split, the output (ytrain/ytest) only contains some of labels' class not all classes. I think it is because the method just split 30% of the whole data but don't care about their labels classes. What should I do to deal with it, I want to make sure the output contain all the labels' classes, and every classes are splited 30% into the train set. Does there have a function can do this? Or I need to reshape my data according to different labels' classes.


1 Answer 1


You can use the argument stratify=Y_source to maintain the proportions after splitting.

  • $\begingroup$ Thanks! it works. $\endgroup$
    – yuyang sun
    Commented Nov 16, 2022 at 16:42
  • $\begingroup$ Glad to help! You can accept the answer if that solved your problem. $\endgroup$
    – db_max
    Commented Nov 16, 2022 at 16:47

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.