如何在Keras多GPU训练中正确使用图片类TF Dataset并解决分片警告

核心原因 #

你设置的分片策略未生效，tf.data.Dataset.with_options() 不会修改原数据集对象，而是返回一个应用了新配置的新数据集对象，你之前的代码没有接收这个返回值，等于配置完全没有生效。
同时这个警告的本质是：默认自动分片会优先尝试文件级分片（每个GPU读取不同的源文件），但你使用from_tensor_slices从内存张量构建数据集，不存在文件级分片的前提，所以默认会自动降级为数据级分片，这个降级行为本身不会影响训练正确性，只是会触发提示。

正确配置步骤 #

1. 数据集构建时应用分片配置，必须接收with_options的返回值：

@tf.function
def load_images(image_path, label):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = vgg16.preprocess_input(image)
    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
    return (image, label)

# 定义数据集配置
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA

training_dataset = tf.data.Dataset.from_tensor_slices((train_paths, train_labels_le))
training_dataset = (
    training_dataset.shuffle(1024)
    .map(load_images, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(BATCH_SIZE) # 此处BATCH_SIZE为全局批次大小，等于单卡批次 * GPU数量即可
    .prefetch(tf.data.AUTOTUNE)
    .with_options(options) # 必须将返回的新数据集赋值给变量
)

2. 训练时移除多余的batch_size参数：
   你已经对数据集做了分批操作，fit方法传入的batch_size会被自动忽略，直接删除即可：

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    keras_model = build_model()
    keras_model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
# 不需要传入batch_size参数
keras_model.fit(training_dataset, epochs=EPOCHS)

可选配置 #

如果你不需要自动分片逻辑，也可以将分片策略设置为tf.data.experimental.AutoShardPolicy.OFF，手动保证全局批次大小可被GPU数量整除即可，同样可以消除警告。

内容的提问来源于stack exchange，提问作者Matthias