tensorflow2.0 - 如何使用 tf.data 同时预处理整个批次？（尤其是在多 GPU 分布式训练中）

Question

我有一个加载图像的 tf.data 管道，目前我在 tf.data.Dataset 上使用单个 GPU 时通过获取单个批次来调整整个批次的大小。但是，我希望在分布式培训中做到这一点，我该怎么做？tf.data.Dataset 映射函数采用单个图像进行预处理，但是我希望做的是在批处理级别执行此操作。

train_dataset = iter(tf.data.Dataset(...))

img,labels = next(train_dataset)
                
scale = random.choice(scales)
method = random.choice(methods)
img = tf.image.resize(img,(scale,scale),method=method)
labels = tf.image.resize(labels,(scale//2,scale//2),method=method)
(img_size//2,img_size//2),method=method)