I am performing multi-class classification problem of different concentrations of Acetaminophen in a specified dataset. My data is in the form of images and I am using CNN. I have compiled all the image paths (with the images being the features) and their respective concentrations (classes) into one excel. Next, I formed a variable "X" to contain all the dataset image arrays and y to contain all the labels, as shown in the code below:

img_path = pd.read_excel(os.path.join(Dir,"DPV Images.xlsx"))

X = []
y = []

for index, row in img_path.iterrows():
    img_array = cv2.imread(row[0]) # Reads the image
    img_class = row[1] # Contains the label

My question is: How can I do a train-test split in this case, to ensure that the train and test data is split evenly between all classes. For example, using scikitlearn's train_test_split with test size of 0.2 will divide the entire dataset into 80% training and 20% training, but it does not ensure that each class is divided into 80% training and 20% testing. How can I go about this?

  • $\begingroup$ Why does the split have to be even? $\endgroup$
    – Dave
    Commented Feb 11, 2023 at 10:35
  • $\begingroup$ Maybe you're looking for stratified sampling? If so you can simply use the stratify parameter. $\endgroup$
    – Erwan
    Commented Feb 11, 2023 at 11:13
  • $\begingroup$ @Dave Because how it currently is I am having random train-test sample proportions for each class (say 1 test sample for class A and the rest is for training, 3 test samples for class B and the rest is for training, etc, noting that the number of samples for each class is the same). I fear this would affect the accuracy of the training/testing, wouldn't it? $\endgroup$
    – Zelreedy
    Commented Feb 11, 2023 at 14:44
  • $\begingroup$ this solution works stackoverflow.com/questions/45516424/… $\endgroup$
    – Vivek Dani
    Commented Oct 24, 2023 at 7:38

2 Answers 2


To do the train-test split in a method that assures an equal distribution of classes between the training and testing sets, utilize the StratifiedShuffleSplit class fromscikit-learn model selection module.


from sklearn.model_selection import StratifiedShuffleSplit
splitter = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=0)
for train_idx, test_idx in splitter.split(X, y):
    X_train = [X[i] for i in train_idx]
    y_train = [y[i] for i in train_idx]
    X_test = [X[i] for i in test_idx]
    y_test = [y[i] for i in test_idx]

You can write a custom function to do this job, below is an example:

import numpy as np
from sklearn.model_selection import train_test_split

def split_data(X, test_size=0.25, random_state=2):
    unique_classes = np.unique(y)
    X_train = []
    X_test = []
    y_train = []
    y_test = []
    for class_value in unique_classes:
        X_class = X[y == class_value]
        y_class = y[y == class_value]
        X_class_train, X_class_test, y_class_train, y_class_test = train_test_split(X_class, y_class, test_size=test_size, random_state=random_state)
    X_train = np.concatenate(X_train)
    X_test = np.concatenate(X_test)
    y_train = np.concatenate(y_train)
    y_test = np.concatenate(y_test)
    return X_train, X_test, y_train, y_test

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.