Scikit-optimize: How can I parallelize a `gp_minimize` over several GPUs?

[frederikfaye]

I have eight GPUs at my disposal for doing a hyperparameter search for a neural network, but for i/o reasons, I would like to not divide minibatches between them.

Until now, I've made eight separate, but identical gp_minimize calls, one for each GPU. This creates eight distinct res.models objects. However, it would clearly be better if all eight GPUs could contribute to create a single res, with a single models object (that would then have eight times the information available to make a decision about where to sample next).

How can I make this possible? Ideas for hacks and quick fixes are appreciated if no canonical way exists.

[betatim]

I would use https://scikit-optimize.github.io/#skopt.Optimizer directly (this is the workhorse behind gp_minize) and ask() for eight points, evaluate each on one of your GPUs and iterate that way.

[frederikfaye]

Thanks for the suggestion, this seems to be exactly what I was looking for.

However, I'm a bit unsure of the implementation. ask() samples from one's search space, but how do I actually give the score found from evaluating a sampled point to my Optimizer instance? A minimal example would be awesome.

[betatim]

Check out the docstrings for the Optimizer class. ask() gives you points to evaluate and tell() let's you give them back to the optimizer (with the result of evaluating them). You could also look at the code of base_minimize() which is what gp_minimize uses as all our minimize functions use Optimzier.

https://github.com/scikit-optimize/scikit-optimize/blob/master/examples/ask-and-tell.ipynb is a minimal example.

[frederikfaye]

@betatim Thanks!

A practical question: (EDIT: Solved in the next comment!)

Let's say I have a list of my idle GPUs, idle_ids, a search space space, as well as a dict scores for keeping track of the scores returned by evaluating on a given GPU id.
How would I distribute my sampled points? If I do

import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'

import skopt

idle_ids = # ...
scores = # ...

def objective():
    # ...
    return score

opt = skopt.Optimizer(space)

sampled_points = opt.ask(len(idle_ids))

for point,id in zip(sampled_points,idle_ids):
    os.environ['CUDA_VISIBLE_DEVICES'] = id
    scores[id] = objective(point)
    # ...

then point[1] won't begin being evaluated before point[0] is done being evaluated, which of course defeats the whole purpose of parallelizing. Could I use threading, or perhaps something involving multiprocessing.pool.AsyncResult? I unfortunately have no experience with threading/multiprocessing, so I don't really know if something like this was what you had in mind.

[frederikfaye]

I think I've solved the parallelization problem:

# Imports
import os
import sys
import time
import pickle
from multiprocessing import Pool
from skopt import Optimizer

# User input
gpu_ids = '0,1,5'
n_jobs = 4
n_calls = 100
n_initial_points = 10
search_space = # [...]

# Define objective function
def objective(point, gpu_id='0'):
    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
    # ...
    return point, score

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids

# Prepare for parallelization over the chosen GPUs
gpu_id_list = gpu_ids.split(',')

# For each GPU, keep track of its results, which has a ready() method,
# telling if it has completed its given task
gpu_state = {id:None for id in gpu_id_list}

# Initialize optimizer, which will keep track of the results received from
# all GPUs
opt = Optimizer(search_space, 
    n_initial_points=n_initial_points, 
    acq_optimizer_kwargs={'n_jobs':n_jobs})

# Begin the process of as many workers as there are visible GPUs
with Pool(processes=len(gpu_id_list)) as pool:

    n_completed_calls = 0
    ready_ids = [id for id in gpu_state]

    while n_completed_calls < n_calls:

        if len(ready_ids) > 0:

            # Update optimizer state
            n_updated = 0
            for id in ready_ids:
                if gpu_state[id] is not None:
                    opt.tell(*gpu_state[id].get())
                    n_updated += 1
            n_completed_calls += n_updated

            # Save optimizer state
            if n_updated > 0:
                with open('optimizer.pkl', 'wb') as f:
                    pickle.dump(opt, f)

            # Sample points for all idle GPUs
            sampled_points = opt.ask(len(ready_ids))

            # Distribute and evaluate the sampled points over the workers
            for point,id in zip(sampled_points,ready_ids):
                gpu_state[id] = pool.apply_async(objective, args=(point,id,))

            # Check if there are any ready (idle) GPUs
            ready_ids = [id for id in gpu_state if gpu_state[id].ready()]

        else:
            # Wait two seconds, so the while-condition isn't checked constantly
            time.sleep(2)

            # Check if there are any ready (idle) GPUs
            ready_ids = [id for id in gpu_state if gpu_state[id].ready()]

I hope this helps someone else wanting the same functionality. Disclaimer: I may have forgotten to define some variable(s), as the above is obviously not a runnable script, but has been composed from snippets of my actual code.