Produce balanced mini batch with Dataset API

Question

I've a question about the new dataset API (tensorflow 1.4rc1). I've a unbalanced dataset wrt to labels 0 and 1. My goal is to create balanced mini batches during the preprocessing.

Assume I've two filtered datasets:

ds_pos = dataset.filter(lambda l, x, y, z: tf.reshape(tf.equal(l, 1), []))
ds_neg = dataset.filter(lambda l, x, y, z: tf.reshape(tf.equal(l, 0), [])).repeat()

Is there a way to combine these two datasets such that the resulting dataset looks like ds = [0, 1, 0, 1, 0, 1]:

Something like this:

dataset = tf.data.Dataset.zip((ds_pos, ds_neg))
dataset = dataset.apply(...)
# dataset looks like [0, 1, 0, 1, 0, 1, ...]
dataset = dataset.batch(20)

My current approach is:

def _concat(x, y):
   return tf.cond(tf.random_uniform(()) > 0.5, lambda: x, lambda: y)
dataset = tf.data.Dataset.zip((ds_pos, ds_neg))
dataset = dataset.map(_concat)

But I've the feeling there is a more elegant way.

Thanks in advance!

Answer 1

You are on the right track. The following example uses Dataset.flat_map() to turn each pair of a positive example and a negative example into two consecutive examples in the result:

dataset = tf.data.Dataset.zip((ds_pos, ds_neg))

# Each input element will be converted into a two-element `Dataset` using
# `Dataset.from_tensors()` and `Dataset.concatenate()`, then `Dataset.flat_map()`
# will flatten the resulting `Dataset`s into a single `Dataset`.
dataset = dataset.flat_map(
    lambda ex_pos, ex_neg: tf.data.Dataset.from_tensors(ex_pos).concatenate(
        tf.data.Dataset.from_tensors(ex_neg)))

dataset = dataset.batch(20)