如何使用Tensorflow进行分布式预测/推断

2024-03-19 07:06:46 发布

您现在位置:Python中文网/ 问答频道 /正文

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

目前,在运行许多大型预测时,我的一个GPU(12gb)的内存(需要17gb),推断失败,因为内存不足:

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

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


Tags: of内存模型镜像gpu分布式集群out
1条回答
网友
1楼 · 发布于 2024-03-19 07:06:46

我能够将单个工人、多个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小时,因此即使是2倍的加速,也不会有8倍的GPU数量

我认为这主要是一个数据传输问题,但我不知道如何解决它。希望其他人能提供一些更好的见解

相关问题 更多 >