How to split data on balanced training set and test set on sklearn

machine-learning scikit-learn svm cross-validation
Jeanne picture Jeanne · Feb 18, 2016 · Viewed 42.5k times · Source

I am using sklearn for multi-classification task. I need to split alldata into train_set and test_set. I want to take randomly the same sample number from each class. Actually, I amusing this function

X_train, X_test, y_train, y_test = cross_validation.train_test_split(Data, Target, test_size=0.3, random_state=0)

but it gives unbalanced dataset! Any suggestion.

Answer

Guiem Bosch picture Guiem Bosch · Feb 18, 2016

Although Christian's suggestion is correct, technically train_test_split should give you stratified results by using the stratify param.

So you could do:

X_train, X_test, y_train, y_test = cross_validation.train_test_split(Data, Target, test_size=0.3, random_state=0, stratify=Target)

The trick here is that it starts from version 0.17 in sklearn.

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

Related questions

How to interpret scikit's learn confusion matrix and classification report?

I have a sentiment analysis task, for this Im using this corpus the opinions have 5 classes (very neg, neg, neu, pos, very pos), from 1 to 5. So I do the classification as follows: from sklearn.feature_extraction.text import TfidfVectorizer import …

Read More
Python : How to find Accuracy Result in SVM Text Classifier Algorithm for Multilabel Class

I have used following set of code: And I need to check accuracy of X_train and X_test The following code works for me in my classification problem over multi-labeled class import numpy as np from sklearn.pipeline import …

Read More
Determining the most contributing features for SVM classifier in sklearn

I have a dataset and I want to train my model on that data. After training, I need to know the features that are major contributors in the classification for a SVM classifier. There is something called feature importance for …

Read More