如何使用tf.Dataset将数据加载到多个GPU？

import tensorflow as tf

# 创建一个tf.Dataset对象，用于加载和预处理数据
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.shuffle(1000).batch(64)

# 使用MirroredStrategy指定多GPU训练策略
strategy = tf.distribute.MirroredStrategy()

# 在strategy.scope()上下文中定义模型和训练过程
with strategy.scope():
    model = tf.keras.Sequential([...])  # 定义模型结构
    optimizer = tf.keras.optimizers.Adam()  # 定义优化器
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy()  # 定义损失函数

    # 定义训练步骤
    def train_step(inputs):
        images, labels = inputs

        with tf.GradientTape() as tape:
            predictions = model(images, training=True)
            loss = loss_object(labels, predictions)

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        return loss

    # 定义分布式训练过程
    @tf.function
    def distributed_train_step(inputs):
        per_replica_losses = strategy.experimental_run_v2(train_step, args=(inputs,))
        return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

    # 执行训练过程
    for epoch in range(num_epochs):
        total_loss = 0.0
        num_batches = 0

        for inputs in dataset:
            total_loss += distributed_train_step(inputs)
            num_batches += 1

        average_loss = total_loss / num_batches
        print("Epoch {}: loss = {}".format(epoch, average_loss))