Stratified train/validation/test split without scikit-learn

Question

I'm working on the mnist dataset with 1797 images indicating 0 to 10 digits. I want to split the dataset into the train, validation, and test sub-data so that the same number of each digit is specified for each sub_data. How can I do stratification without sklearn library in python?

Thank you for your answer in advance.

yann ziselman yann ziselman · Accepted Answer · 2021-05-04T09:00:51

To perform stratified data splitting, you need to know which class each data point belongs to. If you have a list of data points and a corresponding list of classes, you can extract all the points that belong to a certain class and split them according to the input proportions.

Here's some code that implements the idea: Note that you'll have to add some array that tracks the classes the data points belong to after being split in the loop.

import numpy as np
train, valid, test = 0.6, 0.2, 0.2
data_points = np.random.rand(1000, 32, 32)
classes     = np.random.randint(0, 10, size = (1000,))
class_set   = np.unique(classes)
data_train  = []
data_valid  = []
data_test   = []
for class_i in class_set:
    data_inds    = np.where(classes==class_i)
    data_i       = data_points[data_inds, ...]
    N_i          = len(data_inds)
    N_i_train    = int(N_i*train)
    N_i_valid    = int(N_i*valid)
    data_train.append(data_i[:N_i_train])
    data_valid.append(data_i[N_i_train:N_i_train+N_i_valid])
    data_test.append(data_i[N_i_train+N_i_valid:])
    
data_train = np.concatenate(data_train)
data_valid = np.concatenate(data_valid)
data_test = np.concatenate(data_test)

Stratified train/validation/test split without scikit-learn

1 Answers