Skip to content

Instantly share code, notes, and snippets.

@ArthurCamara

ArthurCamara/free_gpus.py

Last active September 16, 2020 09:18
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 ArthurCamara/4296d678a305324aecd624354a27fd58 to your computer and use it in GitHub Desktop.
Save ArthurCamara/4296d678a305324aecd624354a27fd58 to your computer and use it in GitHub Desktop.
Download ZIP
Find free GPUs and set them as visible for pytorch or tensorflow
Raw
free_gpus.py
import os
import numpy as np
def get_free_gpus(n_gpus=2):
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_available = np.array([int(x.split()[2]) for x in open('tmp', 'r').readlines()])
good_gpus = list((np.array(memory_available) > 6000).nonzero()[0])
if len(good_gpus) == 0:
raise IndexError
good_gpus = good_gpus[:n_gpus]
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, good_gpus))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment