Пусть ray автоматически назначит графический процессор для каждой задачи

Question

Я использую ray (https://ray.readthedocs.io/en/latest/)) для запуска задачи, которая подает изображения в сеть Inception v3, загруженную из tenorflow.keras. Моя задача выглядит следующим образом (упрощенно):

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 графических процессоров, но у меня нет луча, чтобы назначить задачу всем графическим процессорам. Если я установил

ray.init(num_gpus=8)

, кажется, что она не назначает 1 задачу на графический процессор.

Мой вопрос: как я могу заставить ray автоматически назначать одну задачу для каждого графического процессора, то есть распараллеливать предсказание для всех моих изображений по всем доступным графическим процессорам?

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