5
$\begingroup$

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 :

enter image description here

$\endgroup$
1
  • $\begingroup$ how you create data distrbution using hostogram? could you share the code here? $\endgroup$ – donto Dec 17 '18 at 10:19
5
$\begingroup$

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!

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.