将 tf.Dataset 与另一个 tf.Dataset 随机交错

Question

I have two datasets:

main_ds = tf.data.Dataset.from_tensor_slices(list(range(1000, 1100)))
backgroud_ds = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])

I want a batch interleaving main_ds and backgroud_ds data randomly. For instance, a batch of size 10 should look like:

[3, 1017, 1039, 3, 2, 1024, 4, 1, 1053, 4]

I tried the following:

def interlace_background(image, background):
    return  tf.cond(tf.random_uniform([]) < .5, lambda: image, lambda: background)

background_ds = background_ds.shuffle(10).repeat(-1)
background_it = background_ds.make_initializable_iterator()
background_next = background_it.get_next()

main_ds = main_ds.shuffle(10)\
                 .repeat(-1)\
                 .map(lambda x: interlace_background(x, background_next))\
                 .batch(10)
main_it = main_ds.make_initializable_iterator()
main_next = main_it.get_next()

but I get a fixed background across all batches:

batch 0: [   3 1006    3 1001    3 1005 1015 1000    3    3]
batch 1: [1007    3 1012 1018 1013    3 1008 1019    3    3]
batch 2: [1016    3 1025    3    3    3 1021    3    3 1035]
batch 3: [1038    3    3 1023 1020    3    3 1046 1034 1047]
batch 4: [   3    3 1039    3    3    3    3    3 1053    3]

Why is the background fixed (cf. above where background is always 3 ) and how could I solve this?

Fully reproducible code below:

import tensorflow as tf
import numpy as np

def interlace_background(image, background):
    return  tf.cond(tf.random_uniform([]) < .5, lambda: image, lambda: background)

main_ds = tf.data.Dataset.from_tensor_slices(list(range(1000, 1100)))
background_ds = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])

background_ds = background_ds.shuffle(10).repeat(-1)
background_it = background_ds.make_initializable_iterator()
background_next = background_it.get_next()

main_ds = main_ds.shuffle(10)\
                 .repeat(-1)\
                 .map(lambda x: interlace_background(x, background_next))\
                 .batch(10)
main_it = main_ds.make_initializable_iterator()
main_next = main_it.get_next()

with tf.Session() as sess:
    sess.run(background_it.initializer)
    sess.run(main_it.initializer)
    for i in range(5):
        print('batch %i' % i, sess.run(main_next))

Answer 1

You can do the same thing with Dataset.zip() and Dataset.map() .

Here is the code:

import tensorflow as tf

def interlace_background(image, background):
    return tf.cond(tf.random_uniform([]) < .5, lambda: image, lambda: background)


main_ds = tf.data.Dataset.from_tensor_slices(list(range(1000, 1100))).shuffle(100)
background_ds = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4]).shuffle(4)

new_ds = tf.data.Dataset \
    .zip((main_ds, background_ds)) \
    .repeat(-1) \
    .map(lambda x, y: interlace_background(x, y)) \
    .batch(10)

iterator = new_ds.make_initializable_iterator()
next_item = iterator.get_next()

with tf.Session() as sess:
    sess.run(iterator.initializer)
    for i in range(5):
        print('batch %i' % i, sess.run(next_item))

Output:

batch 0 [1065    2    4    1    2    4    1 1036 1072 1020]
batch 1 [   4    3    2 1057    1    4    2 1077    3    1]
batch 2 [   3 1044 1042 1049 1029    1    3 1069 1018    3]
batch 3 [   2    4 1089 1094    2 1022 1041 1006    1    3]
batch 4 [1079    2    1    3 1023 1042    4 1018 1054    4]