How to avoid covariate shift in python and distribute classes in each train and test phase?

Question

We all know that with the use of sklearn package from python, we can create X_train, X_test, y_train and y_test via this code:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)

I want to make sure that each training and testing phase of a multi-class data set, have 66.33/33.33 percent of each class values so the prediction and accuracy would get better. All i want is 66.33 percent of class A in training set and 33.33 percent of Class A in test set. And, so on for other classes, like B, C, D and etc. in a given multi-class data set.

Is the code provided enough to achieve this or should i write extra code?

Thanks

Answer 1

You need to call train_test_split(X, y, test_size=0.33, stratify=y) assuming that y contains your class labels. Alternatively, you can call StratifiedShuffleSplit directly.