Stratify batch in Tensorflow 2

Question

I have minibatches that I get from an sqlite database with data of integer and float type, x, and a binary label in 0 and 1, y. I am looking for something like X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(y, x, test_size=0.1, random_state=1, stratify=True) from scikit-learn, where a keyword could stratify the data (i.e. the same number of class-0 and class-1 instances).

In Tensorflow 2, stratification seems not straightforwardly possible. My very complicated solution works for me, but takes a lot of time because of all the reshaping and transposing:

def stratify(x, y):
    # number of positive instances (the smaller class)
    pos = np.sum(y).item() # how many positive bonds there are
    x = np.transpose(x)

    # number of features 
    f = np.shape(x)[1] 

    # filter only class 1
    y = tf.transpose(y)
    x_pos = tf.boolean_mask(x, 
    y_pos = tf.boolean_mask(y, y)

    # filter only class 1
    x_neg = tf.boolean_mask(x, tf.bitwise.invert(y)-254)
    x_neg = tf.reshape(x_neg, [f,-1])
    y_neg = tf.boolean_mask(y, tf.bitwise.invert(y)-254)

    # just take randomy as many class-0 as there are class-1 
    x_neg = tf.transpose(tf.random.shuffle(tf.transpose(x_neg)))
    x_neg = x_neg[:,0:pos]
    y_neg = y_neg[0:pos]

    # concat the class-1 and class-0 together, then shuffle, and concat back together
    x = tf.concat([x_pos,tf.transpose(x_neg)],0)
    y = tf.concat([y_pos, tf.transpose(y_neg)],0)
    xy = tf.concat([tf.transpose(x), tf.cast(np.reshape(y,[1, -1]), tf.float64)],0)
    xy = tf.transpose((tf.random.shuffle(tf.transpose(xy)))) # because there is no axis arg in shuffle
    x = xy[0:f,:]
    x = tf.transpose(x)
    y = xy[f,:]

    return x, y

I am happy to see some feedback/improvement on my own function or novel, easier ideas.

Answer 1

Data division is best if it is done in raw format only or before you transform it into tensors. If there is a strong requirement to do it in TensorFlow only, then I will suggest you to make use of tf.data.Dataset class. I have added the demo code with relevant comments explaining the steps.

import tensorflow as tf
import numpy as np

TEST_SIZE = 0.1
DATA_SIZE = 1000

# Create data
X_data = np.random.rand(DATA_SIZE, 28, 28, 1)
y_data = np.random.randint(0, 2, [DATA_SIZE])
samples1 = np.sum(y_data)
print('Percentage of 1 = ', samples1 / len(y_data))

# Create TensorFlow dataset
dataset = tf.data.Dataset.from_tensor_slices((X_data, y_data))

# Gather data with 0 and 1 labels separately
class0_dataset = dataset.filter(lambda x, y: y == 0)
class1_dataset = dataset.filter(lambda x, y: y == 1)

# Shuffle them
class0_dataset = class0_dataset.shuffle(DATA_SIZE)
class1_dataset = class1_dataset.shuffle(DATA_SIZE)

# Split them
class0_test_samples_len = int((DATA_SIZE - samples1) * TEST_SIZE)
class0_test = class0_dataset.take(class0_test_samples_len)
class0_train = class0_dataset.skip(class0_test_samples_len)

class1_test_samples_len = int(samples1 * TEST_SIZE)
class1_test = class1_dataset.take(class1_test_samples_len)
class1_train = class1_dataset.skip(class1_test_samples_len)

print('Train Class 0 = ', len(list(class0_train)), ' Class 1 = ', len(list(class1_train)))
print('Test Class 0 = ', len(list(class0_test)), ' Class 1 = ', len(list(class1_test)))

# Gather datasets
train_dataset = class0_train.concatenate(class1_train).shuffle(DATA_SIZE)
test_dataset = class0_test.concatenate(class1_test).shuffle(DATA_SIZE)

print('Train dataset size = ', len(list(train_dataset)))
print('Test dataset size = ', len(list(test_dataset)))

Sample output:

Percentage of 1 =  0.474
Train Class 0 =  474  Class 1 =  427
Test Class 0 =  52  Class 1 =  47
Train dataset size =  901
Test dataset size =  99