# Meaning of stratify parameter

I'm training a Neural Network and I'm trying to divide my data into training and testing sets. I have a lot of output classes and for some of them I have as little as 2 examples, so I would like to have, in that case, 1 example in training and 1 example in testing. From what I've read, this is using the stratify parameter, but what does stratify mean?

I'm divifing my data into training and testing:

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.1, random_state=42, stratify=y)


So, from my understanding, this divides into two sets: training (90% of data) and testing (10% of data) but making sure that there are at least 1 of each class in each set?

stratify parameter will preserve the proportion of target as in original dataset, in the train and test datasets as well.

So if your original dataset df has target/label as [0,1,2] in the ratio say, 40:30:30. That is, for every 100 datasets, you can find 40, 30 and 30 observations of target 0,1 and 2 respectively.

Now when you split this original using the train_test_split(x,y,test_size=0.1,stratify=y), the methods returns train and test datasets in the ratio of 90:10. Now in each of these datasets, the target/label data proportion is preserved as 40:30:30 for the classes [0,1,2].

Often, we want to preserve the dataset proportions for better prediction and reproduceability of results

The stratify parameter asks whether you want to retain the same proportion of classes in the train and test sets that are found in the entire original dataset.

For example, if there are 100 observations in the entire original dataset of which 80 are class $$a$$ and 20 are class $$b$$ and you set stratify = True, with a .7 : .3 train-test split, you will get a training set with 56 examples of class $$a$$ and 14 examples of class $$b$$.

In simple word, it ensures that the splitted data have at least some similarity between the train and test data. Suppose a scenario in which you have 100 binary(0,1) datasets of a cancer patient and you have splitted it into 80% as train_data and 20% as test_data. And by default in test_data, you get 15 values as " 1 " and 5 values as 0, So it will be a quite biased data_set. In order to avoid this situation, we use "stratify" to provide a set of common relationship between training and testing dataset.(example: 11 values as "1" and 9 values a "0")