Есть ли способ передать аргументы нескольким заданиям в optuna?

Question

Я пытаюсь использовать optuna для поиска пространств гиперпараметров.

В одном конкретном сценарии я обучаю модель на машине с несколькими графическими процессорами. Модель и размер пакета позволяют мне проводить 1 тренировку на 1 GPU. Итак, в идеале я хотел бы позволить optuna распределять все испытания по доступным графическим процессорам, чтобы на каждом графическом процессоре всегда выполнялась одна пробная версия.

В документации говорится, что я должен просто запустить один процесс для каждого графического процессора в отдельном терминале, например:

CUDA_VISIBLE_DEVICES=0 optuna study optimize foo.py objective --study foo --storage sqlite:///example.db

Я хочу избежать этого, потому что после этого весь поиск гиперпараметров продолжается в несколько раундов. Я не хочу всегда вручную запускать процесс для каждого графического процессора, проверять, когда все будет завершено, а затем начинать следующий раунд.

Я видел study.optimize имеет n_jobsаргумент. На первый взгляд это кажется идеальным.Например, я мог бы сделать это:

import optuna

def objective(trial):
    # the actual model would be trained here
    # the trainer here would need to know which GPU
    # it should be using
    best_val_loss = trainer(**trial.params)
    return best_val_loss

study = optuna.create_study()
study.optimize(objective, n_trials=100, n_jobs=8)

Это запускает несколько потоков, каждый из которых запускает обучение. Однако тренер внутриobjectiveкаким-то образом нужно знать, какой графический процессор он должен использовать. Есть ли уловка для этого?

Answer 1

После нескольких психических срывов я понял, что могу делать то, что хочу, используя multiprocessing.Queue. Чтобы включить его в целевую функцию, мне нужно определить его как лямбда-функцию или как класс (я думаю, что частичное также работает). Например

from contextlib import contextmanager
import multiprocessing
N_GPUS = 2

class GpuQueue:

    def __init__(self):
        self.queue = multiprocessing.Manager().Queue()
        all_idxs = list(range(N_GPUS)) if N_GPUS > 0 else [None]
        for idx in all_idxs:
            self.queue.put(idx)

    @contextmanager
    def one_gpu_per_process(self):
        current_idx = self.queue.get()
        yield current_idx
        self.queue.put(current_idx)


class Objective:

    def __init__(self, gpu_queue: GpuQueue):
        self.gpu_queue = gpu_queue

    def __call__(self, trial: Trial):
        with self.gpu_queue.one_gpu_per_process() as gpu_i:
            best_val_loss = trainer(**trial.params, gpu=gpu_i)
            return best_val_loss

if __name__ == '__main__':
    study = optuna.create_study()
    study.optimize(Objective(GpuQueue()), n_trials=100, n_jobs=8)

Answer 2

Чтобы решить эту проблему, можно ввести глобальную переменную, которая отслеживает, какой GPU используется в данный момент, что затем можно считать в целевой функции. Код выглядит так.

      EPOCHS = n
USED_DEVICES = []

def objective(trial):
    
    time.sleep(random.uniform(0, 2)) #used because all n_jobs start at the same time
    gpu_list = list(range(torch.cuda.device_count())
    unused_gpus = [x for x in gpu_list if x not in USED_DEVICES]
    idx = random.choice(unused_gpus)
    USED_DEVICES.append(idx)
    unused_gpus.remove(idx)
    
    DEVICE = f"cuda:{idx}"
    
    model = define_model(trial).to(DEVICE)
    
    #... YOUR CODE ...

    for epoch in range(EPOCHS):
        
        # ... YOUR CODE ...
            
        if trial.should_prune():
            USED_DEVICES.remove(idx)
            raise optuna.exceptions.TrialPruned()
    
    #remove idx from list to reuse in next trial
    USED_DEVICES.remove(idx)

Answer 3

Если вам нужно задокументированное решение передачи аргументов целевым функциям, используемым несколькими заданиями, то в документации Optuna представлены два решения:

• вызываемые классы (можно комбинировать с многопроцессорностью),
• Обертка лямбда-функции (осторожно: проще, но не работает с многопроцессорностью).

Если вы готовы использовать несколько сокращений, вы можете пропустить некоторые шаблоны, передав глобальные значения (такие константы, как количество используемых графических процессоров) напрямую (через среду Python) в __call__() метод (а не как аргументы __init__()).

Решение вызываемых классов было протестировано на работу (в optuna==2.0.0) с двумя бэкэндами многопроцессорной обработки (локальная / многопроцессорная) и удаленными бэкэндами баз данных (mariadb / postgresql).