Skip to content

Instantly share code, notes, and snippets.

@DmitryUlyanov
Last active July 2, 2021 10:50
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save DmitryUlyanov/a5c37f08dcf0e242a50bf390c176daae to your computer and use it in GitHub Desktop.
Save DmitryUlyanov/a5c37f08dcf0e242a50bf390c176daae to your computer and use it in GitHub Desktop.
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