python - 如何使用 Tensorflow 进行分布式预测/推理

Question

我想使用 TF 2.0 在我的 GPU 集群上运行分布式预测。我使用 MirroredStrategy 训练了一个用 Keras 制作的 CNN 并保存了它。我可以加载模型并在其上使用 .predict()，但我想知道这是否会使用可用的 GPU 自动进行分布式预测。如果没有，我如何运行分布式预测来加速推理并使用所有可用的 GPU 内存？

目前，当运行许多大型预测时，我超过了我的一个 GPU (12gb) 的内存(需要 17gb)并且推理失败，因为它耗尽了内存:

Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.12GiB

但我有多个 GPU，也想使用它们的内存。谢谢。

Answer 1

我能够将单工、多 GPU 预测拼凑如下(将其视为草图 - 它使用不普遍适用的管道代码，但应该为您提供一个模板):

# https://github.com/tensorflow/tensorflow/issues/37686
# https://www.tensorflow.org/tutorials/distribute/custom_training
def compute_and_write_ious_multi_gpu(path: str, filename_csv: str, include_sampled: bool):
    strategy = tf.distribute.MirroredStrategy()
    util.log('Number of devices: {}'.format(strategy.num_replicas_in_sync))
    (ds, s, n) = dataset(path, shuffle=False, repeat=False, mask_as_input=True)
    dist_ds = strategy.experimental_distribute_dataset(ds)

    def predict_step(inputs):
        images, labels = inputs
        return model(images, training=False)

    @tf.function
    def distributed_predict_step(dataset_inputs):
        per_replica_losses = strategy.run(predict_step, args=(dataset_inputs,))
        return per_replica_losses  # unwrap!?

    # https://stackoverflow.com/questions/57549448/how-to-convert-perreplica-to-tensor
    def unwrap(per_replica):  # -> list of numpy arrays
        if strategy.num_replicas_in_sync > 1:
            out = per_replica.values
        else:
            out = (per_replica,)
        return list(map(lambda x: x.numpy(), out))

    with strategy.scope():
        model = wrap_model()

    util.log(f'Starting distributed prediction for {filename_csv}')
    ious = [unwrap(distributed_predict_step(x)) for x in dist_ds]
    t = ious
    ious = [item for sublist in t for item in
            sublist]  # https://stackoverflow.com/questions/952914/how-to-make-a-flat-list-out-of-list-of-lists
    util.log(f'Distributed prediction done for {filename_csv}')
    ious = np.concatenate(ious).ravel().tolist()
    ious = round_ious(ious)
    ious = list(zip(ious, ds.all_image_paths))
    ious.sort()
    write_ious(ious, filename_csv, include_sampled)

这确实在 GPU 上分配了负载，但不幸的是它们的使用非常差 - 在我的特殊情况下，相应的单 GPU 代码运行时间约为 12 小时，而运行时间为 7.7 小时，因此即使有 8 倍的加速也没有 2 倍的加速GPU 的数量。
我认为这主要是数据馈送问题，但我不知道如何解决。希望其他人可以提供一些更好的见解？