Before training a model, I would like to split the data in a 80/20 ratio.

For example, in my dataset of 3000 rows, I have different labels. Out of which A is one label name, has 100 records in the dataset. How can I divide the dataset so that 80% of the label A data is in training and 20% is in the test set.

Do we have any pre defined functions for this?


1 Answer 1


Use a stratified test split.

import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit
X = np.array([[1, 3], [3, 7], [2, 4], [4, 8]])
y = np.array([0, 1, 0, 1])
stratSplit = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=42)
for train_idx, test_idx in stratSplit:

# [[3 7]
#  [2 4]]
# [1 0]

Source: https://stackoverflow.com/a/35473089/9524424


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.