Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
from __future__ import print_function
import threading
from joblib import Parallel, delayed
import Queue
import os
# Fix print
_print = print
_rlock = threading.RLock()
def print(*args, **kwargs):
with _rlock:
_print(*args, **kwargs)
# Define number of GPUs available
N_GPU = 4
# Put indices in queue
q = Queue.Queue(maxsize=N_GPU)
for i in range(N_GPU):
q.put(i)
def runner(x):
gpu = q.get()
print (x, gpu)
# Put here your job cmd
cmd = "python main.py %s" % x
os.system("CUDA_VISIBLE_DEVICES=%d %s" % (gpu, cmd))
# return gpu id to queue
q.put(gpu)
# Change loop
Parallel(n_jobs=N_GPU, backend="threading")(delayed(runner)(i) for i in range(100))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment