How can l get 50 % examples in training set and 50% in test set for each class when splitting data?

Question

l have a dataset of 200 examples with 10 classes. l would like to split the dataset into training set 50% and test set 50%.

for each class, l have 20 examples. Hence, l would like to get for each class : 10 training examples and 10 test examples.

Here are my classes :

classes=['BenchPress', 'ApplyLipstick', 'BabyCrawling', 'BandMarching', 'Archery', 'Basketball', 'ApplyEyeMakeup', 'BalanceBeam', 'BaseballPitch', 'BasketballDunk']

l tried the following :

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(final_data, true_label, test_size=0.50, random_state=42)

However it returns a 50% training set and 50 % test set, without respecting the proportion for each class (l would like to get 10 examples in test set and 10 examples in training set for each class). Here is the resulted splitting :

Answer 1

sklearn version 0.17 onwards the train_test_split should give you stratified results by using the stratify parameter.

Example code:

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(final_data, true_label, test_size=0.50, random_state=42, stratify=true_label)

From the documentation about the parameter stratify:

stratify: array-like or None (default is None) If not None, data is split in a stratified fashion, using this as the labels array. New in version 0.17: stratify splitting

Hope this helps!