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

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，但我没有得到 ray 将任务分配给所有 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

Let ray automatically assign a GPU per task

python

ray