Question

我正在使用ray（https://ray.readthedocs.io/en/latest/）运行一个任务，该任务将图像馈送到从tensorflow.keras加载的Inception v3网络中。我的任务如下所示（简化）：

from tensorflow.keras.applications import InceptionV3

@ray.remote
def predict(image):
    model = InceptionV3(include_top=False, pooling='avg', input_shape=(1920, 1080, 3))
    return model.predict(image)

我的机器有8个GPU，但是我没有得到将任务分配给所有GPU的机会。如果我设置

ray.init(num_gpus=8)

似乎没有为每个GPU分配1个任务。

我的问题：如何使ray自动为每个GPU分配一项任务，即在所有可用GPU上并行化对我所有图像的预测？

Answer 1


@ray.remote(num_gpus=1)
def predict(image):
    model = InceptionV3(include_top=False, pooling='avg', input_shape=(1920, 1080, 3))
    return model.predict(image)

应该工作。参见https://ray.readthedocs.io/en/latest/using-ray-with-gpus.html#using-remote-functions-with-gpus

让ray为每个任务自动分配GPU

1 个答案: