How to use sklearn train_test_split to stratify data for multi-label classification?

I am attempting to mirror a machine learning program by Ahmed Besbes, but scaled up for multi-label classification. It seems that any attempt to stratify the data returns the following error: The least populated class in y has only 1 member, which is too few. The minimum number of labels for any class cannot be less than 2.

In my data set, I have 1 column which contains clean, tokenized text. The other 8 columns are for the classifications based on the content of that text. Just to note, column 1 - 4 have significantly more samples than 5 - 8 (more obscure classifications derived from the text).

Here is a generic sample from my code:

x = data['cleaned_text']
y = data[['car','truck','ford','chevy','black','white','parked', 'driving']]

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

print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)


Output: (6293,) (700,) (6293, 8) (700, 8)

Adding stratify=y to train_test_split returns the error previously mentioned. Even when I limit y to just one column, I still get the error.

How can I stratify the data so that I give the program a fair look in the training set?

• You try to predict more than one class at the same time. It's not a multi-class classification, but a multi-label classification problem. Please add a sample of your dataset since it is not clear what you try to do. – Tasos Feb 6 '19 at 16:57
• Thanks, I edited the title and body of the initial question to reflect multi-label vice multi-class. As for the data, I can give a generic example – Michael Joy Feb 6 '19 at 17:08

The error you're getting indicates it cannot do a stratified split because one of your classes has only one sample. You need at least two samples of each class in order to put one in the training split and one in the test split. You should examine what your class breakdown is to find the culprit.

Try this:

from skmultilearn.model_selection import iterative_train_test_split X_train, y_train, X_test, y_test = iterative_train_test_split(x, y, test_size = 0.1)

Since you're doing multilabel classification, it's very likely to get unique combinations of each class, which is what causes the error with sklearn. You have to use a special library for multilabel stratified splitting.

More details on how to use skmultilearn

There is a seperate module for classes stratification and no one is going to suggest you to use the train_test_split for this. This could be achieved as follows:

from sklearn.model_selection import StratifiedKFold

train_all = []
evaluate_all = []
skf = StratifiedKFold(n_splits=cv_total, random_state=1234, shuffle=True)
for train_index, evaluate_index in skf.split(train_df.index.values, train_df.coverage_class):
train_all.append(train_index)
evaluate_all.append(evaluate_index)
print(train_index.shape,evaluate_index.shape) # the shape is slightly different in different cv, it's OK

# Getting each batch
def get_cv_data(cv_index):
train_index = train_all[cv_index-1]
evaluate_index = evaluate_all[cv_index-1]
x_train = np.array(train_df.images[train_index].map(upsample).tolist()).reshape(-1, img_size_target, img_size_target, 1)
y_train = np.array(train_df.masks[train_index].map(upsample).tolist()).reshape(-1, img_size_target, img_size_target, 1)
x_valid = np.array(train_df.images[evaluate_index].map(upsample).tolist()).reshape(-1, img_size_target, img_size_target, 1)
y_valid = np.array(train_df.masks[evaluate_index].map(upsample).tolist()).reshape(-1, img_size_target, img_size_target, 1)
return x_train,y_train,x_valid,y_valid

# Training loop
for cv_index in range(cv_total):
x_train, y_train, x_valid, y_valid =  get_cv_data(cv_index+1)
history = model.fit(x_train, y_train,
validation_data=[x_valid, y_valid],
epochs=epochs)



This is a simple code snippet for using StratifiedKFold with your code. Just replace the required parameters and hyper-parameters accordingly.