Is there a way to pass arguments to multiple jobs in optuna?

Question

I am trying to use optuna for searching hyper parameter spaces.

In one particular scenario I train a model on a machine with a few GPUs. The model and batch size allows me to run 1 training per 1 GPU. So, ideally I would like to let optuna spread all trials across the available GPUs so that there is always 1 trial running on each GPU.

In the docs it says, I should just start one process per GPU in a separate terminal like:

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

I want to avoid that because the whole hyper parameter search continues in multiple rounds after that. I don't want to always manually start a process per GPU, check when all are finished, then start the next round.

I saw study.optimize has a n_jobs argument. At first glance this seems to be perfect. E.g. I could do this:

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)

This starts multiple threads each starting a training. However, the trainer within objective somehow needs to know which GPU it should be using. Is there a trick to accomplish that?

score 8 · Accepted Answer · answered May 14 '20 at 14:01

After a few mental breakdowns I figured out that I can do what I want using a multiprocessing.Queue. To get it into the objective function I need to define it as a lambda function or as a class (I guess partial also works). E.g.

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)

score 0 · Answer 2 · answered Aug 29 '20 at 11:53

If you want a documented solution of passing arguments to objective functions used by multiple jobs, then Optuna docs present two solutions:

callable classes (it can be combined with multiprocessing),
lambda function wrapper (caution: simpler, but does not work with multiprocessing).

If you are prepared to take a few shortcuts, then you can skip some boilerplate by passing global values (constants such as number of GPUs used) directly (via python environment) to the __call__() method (rather than as arguments of __init__()).

The callable classes solution was tested to work (in optuna==2.0.0) with the two multiprocessing backends (loky/multiprocessing) and remote database backends (mariadb/postgresql).

yes this is a way to provide initial arguments and I'm using it in the above answer to pass the queue class. But I don't see how this alone can enable me to make jobs assess which GPU is currently free. I think it is still necessary to have something like a queue. — mRcSchwering, Aug 31 '20 at 09:20

score 0 · Answer 3 · edited Sep 08 '22 at 09:32

To overcome the problem if introduced a global variable that tracks, which GPU is currently in use, which can then be read out in the objective function. The code looks like this.

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)

Is there a way to pass arguments to multiple jobs in optuna?

3 Answers3

Linked