2

I am trying to implement Classification algorithm for Iris Dataset (Downloaded from Kaggle). In the Species column the classes (Iris-setosa, Iris-versicolor , Iris-virginica) are in sorted order. How can I stratify the train and test data using Scikit-Learn?

Sarath
  • 163
  • 5
  • 14
  • Does this answer your question? [How to split data into 3 sets (train, validation and test)?](https://stackoverflow.com/questions/38250710/how-to-split-data-into-3-sets-train-validation-and-test) – KJTHoward Mar 04 '20 at 16:34

3 Answers3

7

If you want to shuffle and split your data with 0.3 test ratio, you can use

sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, shuffle=True)

where X is your data, y is corresponding labels, test_size is the percentage of the data that should be held over for testing, shuffle=True shuffles the data before splitting

In order to make sure that the data is equally splitted according to a column, you can give it to the stratify parameter.

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, 
                                                    shuffle=True, 
                                stratify = X['YOUR_COLUMN_LABEL'])
metalrt
  • 356
  • 3
  • 11
2

To make sure that the three classes are represented equally in your train and test, you can use the stratify parameter of the train_test_split function.

from sklearn.model_selection import train_test_split
X_train, y_train, X_test, y_test  = train_test_split(X, y, stratify = y)

This will make sure that the ratio of all the classes is maintained equally.

Vatsal Gupta
  • 471
  • 3
  • 8
-1

use sklearn.model_selection.train_test_split and play around with Shuffle parameter.

shuffle: boolean, optional (default=True) Whether or not to shuffle the data before splitting. If shuffle=False then stratify must be None.

Surbhi Gupta
  • 19
  • 3
  • Shuffle does not equally split the training and test data based on the target. Thanks though! – Sarath Mar 07 '20 at 06:39