Skip to content

Instantly share code, notes, and snippets.

@kris-singh
Forked from DmitryUlyanov/run_batch.py
Created November 8, 2019 04:36
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save kris-singh/5bbd24933089ecad452c04218d278f8d to your computer and use it in GitHub Desktop.
Save kris-singh/5bbd24933089ecad452c04218d278f8d to your computer and use it in GitHub Desktop.
from joblib import Parallel, delayed
import Queue
import os
# 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