I have an extreme multilabel dataset that contains thousands of labels, each label exists at least 10 times.

What is the best way to split the data in a stratified way?

I tried the iterative_train_test_split function from scikit-multilearn, but it didn't work.

Sometimes the kernel crashes, sometimes I get strange errors like KeyError: 'key of type tuple not found and not a MultiIndex'.

I work on a mac with an M1 processor if it changes anything.


  • $\begingroup$ Imho there's no strong reason to use stratified sampling here. In any case you should probably discard rare labels. $\endgroup$
    – Erwan
    Commented Sep 22, 2021 at 16:08
  • $\begingroup$ Thanks. Do you recommend this approach, because you think that if it's so rare in my train data, the probability to see it later on unseen data is very low? $\endgroup$
    – Avihai
    Commented Sep 22, 2021 at 18:20
  • 1
    $\begingroup$ I'm assuming that among thousands of labels there are many which appear very rarely. The ones which appear only one are simply unusable since they can't be in both the training and test set. For the others it depends, but I would suggest trying to start with a subset of the most frequent labels first, then once it works well extend it to take into account other labels. $\endgroup$
    – Erwan
    Commented Sep 22, 2021 at 19:18

2 Answers 2


Usually problems arise if the data contains only one data entry for a given label used for stratification. So before performing stratification drop all the rows with unique labels. You can do this using collections.Counter class. After you dropped those rows it is easy to stratify the dataframe, which you are using I assume, e.g.,

from sklearn import datasets
from sklearn.model_selection import train_test_split

iris = datasets.load_iris()
X = pd.DataFrame(iris.data)
y = iris.target
# these labels will not cause any problems
X['cat'] = np.random.choice(['label1','label2','label3','label4'],len(X))

# but these ones will, because they are unique
X.loc[37, 'cat'] = 'label5'
X.loc[137, 'cat'] = 'label6'

# this row will raise an exception if uncommented
#X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, 
#                                                    random_state=43, stratify = X.cat)

# let's drop rows with unique labels
from collections import Counter
unique_labels = [lab for lab, count in Counter(X.cat).items() if count == 1]

print(f"unique labels to be dropped: {unique_labels}")

# drop rows with unique labels
X = X[~X.cat.isin(unique_labels)]
y = y[X.index]

# now datasets X and y can be used in train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, 
                                                    random_state=43, stratify = X.cat)
  • $\begingroup$ Thanks, but there are no rows with unique labels, so I think that it may be related to the very large number of labels that I have (thousands). $\endgroup$
    – Avihai
    Commented Sep 22, 2021 at 18:21
  • $\begingroup$ @Avihai not at all. I would be happy to help you further, if you provided us with minimally reproducible example, say on colab.research.google.com. $\endgroup$ Commented Sep 23, 2021 at 1:51

1000’s of target labels is going to cause issues with most standard desktop ML algorithm implementations I suspect. Some of the public DL neural networks might be capable of handling this depending on your data and use case.

For a desktop solution, I would suggest breaking your problem down into 1000’s of individual binary classification tasks per target.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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