如何在TensorFlow 2中迭代多个数据集

Question

我使用的是TensorFlow 2.2.0。在我的数据管道中，我使用多个数据集来训练神经网络。比如：

# these are all tf.data.Dataset objects:
paired_data = get_dataset(id=0, repeat=False, shuffle=True)
unpaired_images = get_dataset(id=1, repeat=True, shuffle=True)
unpaired_masks = get_dataset(id=2, repeat=True, shuffle=True)

在训练循环中，我想迭代paired_data来定义一个历元。但我还想迭代unpaired_images和unpaired_masks以优化其他目标（用于语义分割的经典半监督学习，带有掩码鉴别器）

为此，我当前的代码如下所示：

def train_one_epoch(self, writer, step, paired_data, unpaired_images, unpaired_masks):

    unpaired_images = unpaired_images.as_numpy_iterator()
    unpaired_masks = unpaired_masks.as_numpy_iterator()

    for images, labels in paired_data:

        with tf.GradientTape() as sup_tape, \
                tf.GradientTape() as gen_tape, \
                tf.GradientTape() as disc_tape:

            # paired data (supervised cost):
            predictions = segmentor(images, training=True)
            sup_loss = weighted_cross_entropy(predictions, labels)

            # unpaired data (adversarial cost):
            pred_real = discriminator(next(unpaired_masks), training=True)
            pred_fake = discriminator(segmentor(next(unpaired_images), training=True), training=True)
            gen_loss = generator_loss(pred_fake)
            disc_loss = discriminator_loss(pred_real, pred_fake)

        gradients = sup_tape.gradient(sup_loss, self.segmentor.trainable_variables)
        generator_optimizer.apply_gradients(zip(gradients, self.segmentor.trainable_variables))

        gradients = gen_tape.gradient(gen_loss, self.segmentor.trainable_variables)
        generator_optimizer.apply_gradients(zip(gradients, self.segmentor.trainable_variables))

        gradients = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
        discriminator_optimizer.apply_gradients(zip(gradients, self.discriminator.trainable_variables))

但是，这会导致错误：

main.py:275 train_one_epoch  *
        unpaired_images = unpaired_images.as_numpy_iterator()
    /home/venvs/conda/miniconda3/envs/tf-gpu/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py:476 as_numpy_iterator  **
        raise RuntimeError("as_numpy_iterator() is not supported while tracing "

    RuntimeError: as_numpy_iterator() is not supported while tracing functions

你知道这是怎么回事吗？在tensorflow 2中，这是优化多个损失/数据集的正确方法吗

---

我将当前解决方案添加到评论中。任何关于更优化方法的建议都是非常受欢迎的！：）

Answer 1

我当前的解决方案：

def train_one_epoch(self, writer, step, paired_data, unpaired_images, unpaired_masks):

    # create a new dataset zipping the three original dataset objects
    dataset = tf.data.Dataset.zip((paired_data, unpaired_images, unpaired_masks))
    dataset = dataset.batch(1)

    for (images, labels), unpaired_images, unpaired_masks in dataset:

        # access the elements as the first and only element of the batched dataset
        images, labels, unpaired_images, unpaired_masks = \
            images[0], labels[0], unpaired_images[0], unpaired_masks[0]

        # go ahead and train:
        with tf.GradientTape() as tape:
            #[...]