使用Dataset API生成平衡的迷你批次

时间:2017-10-25 17:24:57

标签: tensorflow tensorflow-datasets

我有关于新数据集API的问题(tensorflow 1.4rc1)。 我标记了01的不平衡数据集。我的目标是在预处理过程中创建平衡的迷你批次。

假设我有两个过滤的数据集:

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()

有没有办法合并这两个数据集,使得结果数据集看起来像ds = [0, 1, 0, 1, 0, 1]

这样的事情:

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)

我目前的做法是:

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)

但我觉得有一种更优雅的方式。

提前致谢!

1 个答案:

答案 0 :(得分:4)

你走在正确的轨道上。以下示例使用Dataset.flat_map()将每个正面示例和一个反面示例转换为结果中的两个连续示例:

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)