# How to split train/test datasets having equal classes proportion

I would like to know how I can split in an equal number the following

Target
0    1586
1     318


in order to have the same proportion of 0 and 1 classes in a dataset to train, if my dataset is called df and includes 10 columns, both numerical and categorical.

I would consider the following

y=df['Target']
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, stratify=y)


so to do a stratification, but I do not know if it is right and I would appreciate if you could confirm it or provide an alternative to do that.

Sample of data

Fin                 Eco   Target
High percentage     12      1
Low percentage      5       0
Medium percentage   48      0
NA                  3       1
TBC                 NA      1
Low percentage      25      0
Medium percentage   12      0


How can I check if it is actually splitting in equal classes proportion my dataset? I think the best way to train a model should be having an equal proportion of 0 and 1 values. Right now I have 5 times data with Target=0.

If what you want to do is keep the same proportions across the splits, what you are doing is right.

In order to validate properly your model, the class distribution should be constant along with the different splits (train, validation, test).

In the train test split documentation , you can find the argument:

stratifyarray-like, default=None If not None, data is split in a stratified fashion, using this as the class labels.

One step beyond will be using Stratified K-Folds cross-validator.

This cross-validation object is a variation of KFold that returns stratified folds. The folds are made by preserving the percentage of samples for each class.

There are more splitting techniques in scikit learn that you can use, have a look.

In order to check what is the function doing just calculate the percentages in the splits:

 np.unique(y_train, return_counts=True)
np.unique(y_val, return_counts=True)


But this will make you have the same proportions across the whole data, if your original label proportion is 1/5, then you will have 1/5 in train and 1/5 in test

If what you want is have the same proportion of classes 50% - 0 and 50% - 1. Then there is two techniques oversampling and undersampling. But I wont recommend you this for your problem. Your label seems fairly good balanced 1/5 its a great proportion.

Still this library has some implementations to do that

In this blog you can see an overview of imbalanced datasets but yours is not. Choosing a proper metric is more important.

• I kindly suggest you read the question more closely. OP asks if what they are doing "is actually splitting in equal classes proportion" their dataset; what they need to hear is that stratify=y will not do that (it will simply maintain the proportions existing in y to the produced datasets y_train and y_val), and arguably not that what they are doing is "right". Oct 11, 2020 at 22:29
• Thanks @desertnaut, great appreciation. I edited Oct 12, 2020 at 6:23

you can try stratified sampling method

from sklearn.model_selection import StratifiedShuffleSplit
split=StratifiedShuffleSplit(n_split=1, test_size=0.2, random_state=9)