keras - 如何使用 keras.utils.Sequence 数据生成器和 tf.distribute.MirroredStrategy 在 tensorflow 中进行多 GPU 模型训练？

Question

我想使用 tensorflow 2.0 在多个 GPU 上训练模型。在用于分布式训练的 tensorflow 教程 ( https://www.tensorflow.org/guide/distributed_training ) 中， tf.data datagenerator 转换为分布式数据集如下:

dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

但是，我想改用我自己的自定义数据生成器(例如， keras.utils.Sequence 数据生成器，以及用于异步批处理生成的 keras.utils.data_utils.OrderedEnqueuer)。但是 mirrored_strategy.experimental_distribute_dataset方法仅支持 tf.data数据生成器。我如何改用 keras 数据生成器？

谢谢!

Answer 1

我用过 tf.data.Dataset.from_generator和我的 keras.utils.sequence在同样的情况下，它解决了我的问题!

train_generator = SegmentationMultiGenerator(datasets, folder) # My keras.utils.sequence object

def generator():
    multi_enqueuer = OrderedEnqueuer(train_generator, use_multiprocessing=True)
    multi_enqueuer.start(workers=10, max_queue_size=10)
    while True:
        batch_xs, batch_ys, dset_index = next(multi_enqueuer.get()) # I have three outputs
        yield batch_xs, batch_ys, dset_index

dataset = tf.data.Dataset.from_generator(generator,
                                         output_types=(tf.float64, tf.float64, tf.int64),
                                         output_shapes=(tf.TensorShape([None, None, None, None]),
                                                        tf.TensorShape([None, None, None, None]),
                                                        tf.TensorShape([None, None])))

strategy = tf.distribute.MirroredStrategy()

train_dist_dataset = strategy.experimental_distribute_dataset(dataset)

请注意，这是我的第一个工作解决方案 - 目前我发现将“无”放在我发现有效的实际输出形状的位置最方便。