如何使用 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 的数量。

我认为这主要是一个数据馈送问题，但我不知道如何解决它。希望其他人可以提供一些更好的见解吗？