在 tensorflow 中在多个数据集上联合训练单个模型

Question

我想在多个数据集上联合训练单个变分自动编码器模型甚至标准自动编码器（例如 mnist、cifar、svhn 等，其中数据集中的所有图像都被调整为相同的输入形状）。这是我使用的 tensorflow 中的 VAE 教程：https://www.tensorflow.org/tutorials/generative/cvae。

为了训练模型，我想从我的数据集中采样（选择）一个数据集，然后在训练循环的每个梯度更新步骤中从该数据集中获取一批图像。我可以将所有数据集组合成一个大数据集，但我想利用给定批次中的图像来自与辅助信息相同的数据集（我仍在弄清楚这部分，但细节并不是太重要，因为我的问题集中在数据管道上）。

我不确定如何进行数据管道设置。本教程指定数据集管道如下：

train_dataset = (tf.data.Dataset.from_tensor_slices(train_images)
                 .shuffle(train_size).batch(batch_size))
test_dataset = (tf.data.Dataset.from_tensor_slices(test_images)
                .shuffle(test_size).batch(batch_size))

其中 train_images 和 test_images 是处理后的 MNIST 数据。因此，它创建了一个 tensorflow 数据集，对整个数据集进行打乱，并将数据批量化为大小为 batch_size 的批次。就我而言，我假设我想为我的数据集中的每个数据集创建一个单独的 train_dataset/test_dataset（例如 cifar_train_dataset/cifar_test_dataset、mnist_train_dataset/mnist_test_dataset 等）。

在培训方面，他们指定的程序如下：

for epoch in range(1, epochs + 1):
  for train_x in train_dataset:
    train_step(model, train_x, optimizer)

  loss = tf.keras.metrics.Mean()
  for test_x in test_dataset:
    loss(compute_loss(model, test_x))
  elbo = -loss.result()
  print('Epoch: {}, Test set ELBO: {})

我可以指定训练迭代/步骤的总数（例如 500,000），而不是指定时期。在每个训练步骤中，我希望从一组数据集（假设概率相等）中采样一个数据集，而不是像上面那样假设单个训练数据集。

现在是我不确定的部分。 for train_x in train_dataset 行是一个循环，它分批迭代整个数据集。相反，我只想为我采样的给定数据集获取一批图像，进行模型更新，然后重复该过程。但是，我不确定如上所述指定数据集是否提供了这种灵活性？有什么方法可以索引批次/获取单个批次，而不是遍历所有批次。

总之，我想在进行模型更新时，通过在每个训练步骤从给定数据集中采样一批图像来在多个数据集上训练单个模型。我完全愿意接受其他解决此问题的建议和方法。谢谢！

Answer 1

如果我正确理解您的问题，您希望控制从训练集和测试集中提取的批次数量，而不是在进行更新之前完全迭代它们。您可以通过将数据集包装在 iter() 中并将其转换为迭代器，并使用 next() 方法获取下一批。

示例：

import numpy as np
import tensorflow as tf


# fake mnist data
train_imgs = tf.random.normal([100, 28, 28, 1])
test_imgs = tf.random.normal([100, 28, 28, 1])
train_labels = tf.one_hot(
    tf.random.uniform([100,], minval=0, maxval=10, dtype=tf.int64), 10)
test_labels = tf.one_hot(
    tf.random.uniform([100,], minval=0, maxval=10, dtype=tf.int64), 10)

# create train/test dataset
train_ds = tf.data.Dataset.from_tensor_slices((train_imgs, train_labels))
train_ds = train_ds.repeat().shuffle(1 << 6).batch(8)
test_ds = tf.data.Dataset.from_tensor_slices((test_imgs, train_labels))
test_ds = test_ds.repeat().shuffle(1 << 6).batch(8)

# simple mnist network
x_in = tf.keras.Input((28, 28, 1))
x = tf.keras.layers.Flatten()(x_in)
x = tf.keras.layers.Dense(100)(x)
x_out = tf.keras.layers.Dense(10)(x)

# simple mnist model
model = tf.keras.Model(x_in, x_out)

# make datasets iterators
train_iter = iter(train_ds)
test_iter = iter(test_ds)

# loss
def xent_loss(y_true, y_pred):
    ce = tf.keras.losses.CategoricalCrossentropy()
    return ce(y_true, y_pred)
  
  
# simple training loop where you control the batches per epoch
# for your train and test datasets
NUM_EPOCHS = 10
NUM_TRAIN_BATCHES_PER_EPOCH = 20
NUM_TEST_BATCHES_PER_EPOCH = 5

for epoch in range(NUM_EPOCHS):
    train_losses = []
    # train
    for _ in range(NUM_TRAIN_BATCHES_PER_EPOCH):
        X_train, y_train = next(train_iter)
        y_hat = model(X_train)
        loss = xent_loss(y_train, y_hat)
        train_losses.append(loss)
        # do gradient update ...
        
    # report train loss
    print(f"epoch: {epoch}\ttrain_loss: {np.mean(train_losses):.4f}")
    train_losses = []
    
    # validate
    test_losses = []
    for _ in range(NUM_TEST_BATCHES_PER_EPOCH):
        X_test, y_test = next(test_iter)
        y_hat = model(X_test)
        loss = xent_loss(y_test, y_hat)
        test_losses.append(loss)
        
    # report validation loss
    print(f"epoch: {epoch}\ttest_loss: {np.mean(test_losses):.4f}")
    test_losses = []
    print('-' * 40)

# epoch: 0  train_loss: 7.3092
# epoch: 0  test_loss: 7.3427
# ----------------------------------------
# epoch: 1  train_loss: 6.8050
# epoch: 1  test_loss: 8.4867
# ----------------------------------------