如何在GANs训练中使用TensorFlow Dataset API?

时间:2018-08-22 16:28:54

标签: tensorflow tensorflow-datasets

我正在训练GAN模型。为了加载数据集,我正在使用TensorFlow的数据集API。

# train_dataset has image and label. z_train dataset has noise (z).
train_dataset = tf.data.TFRecordDataset(train_file)
z_train = tf.data.Dataset.from_tensor_slices(tf.random_uniform([total_training_samples, seq_length,  z_dim],
                                                                 minval=0, maxval=1, dtype=tf.float32))

train_dataset = tf.data.Dataset.zip((train_dataset, z_train))

创建迭代器:

iter = tf.data.Iterator.from_structure(train_dataset.output_types, train_dataset.output_shapes)

使用迭代器:

(img, label), z = iter.get_next()
train_init_op = iter.make_initializer(train_dataset)

在会话中训练GAN时:
首先训练鉴别器:

_, disc_loss = sess.run([disc_optim, disc_loss])

然后训练生成器:

_, gen_loss = sess.run([gen_optim, gen_loss])

这就是陷阱。因为我在鉴别器图和生成器图中都使用了 label 作为条件(CGAN),所以使用两个sess.run在同一运行期间会生成两组不同的 label 批次批次。

for epoch in range(num_of_epochs):
    sess.run([tf.global_variables_initializer(), train_init_op.initializer])
    for batch in range(num_of_batches):
        _, disc_loss = sess.run([disc_optim, disc_loss])
        _, gen_loss = sess.run([gen_optim, gen_loss])

由于我必须在生成器的会话运行中提供与鉴别器的会话运行相同的 label 批处理,因此如何防止Dataset API在一个批处理的同一循环中产生两个不同的批处理?
注意:我正在使用TensorFlow v1.9
谢谢!

1 个答案:

答案 0 :(得分:0)

您可以为同一数据集创建2个迭代器。如果需要重新整理数据集,甚至可以通过将种子指定为张量来实现。参见下面的示例。

import tensorflow as tf

seed_ts = tf.placeholder(tf.int64)
ds = tf.data.Dataset.from_tensor_slices([1,2,3,4,5]).shuffle(5, seed=seed_ts, reshuffle_each_iteration=True)
it1 = ds.make_initializable_iterator()
it2 = ds.make_initializable_iterator()

input1 = it1.get_next()
input2 = it2.get_next()

with tf.Session() as sess:
    for ep in range(10):
        sess.run(it1.initializer, feed_dict={seed_ts: ep})
        sess.run(it2.initializer, feed_dict={seed_ts: ep})

        print("Epoch" + str(ep))
        for i in range(5):
            x = sess.run(input1)
            y = sess.run(input2)
            print([x, y])
相关问题