auto-select free CUDA devices for training in python

get_free_gpu.py

import subprocess, sys
import numpy as np
from io import StringIO


def get_free_gpu(num_gpus=1):
    """ returns: comma-separated list of device ids, e.g. to set CUDA_VISIBLE_DEVICES for PyTorch """
    output = subprocess.check_output(["nvidia-smi -q -d Memory | grep -A4 GPU | grep Free"], shell=True).decode()
    memory_available = np.array( [int(x.split()[2]) for x in StringIO(output)] )
    memory_available = np.where( memory_available == 11018 )[0]
    if len(memory_available) >= num_gpus:
        gpus = memory_available[0:num_gpus]
        return ','.join( map(str, gpus) )
    else:
        sys.exit(f"There are no {num_gpus} GPUs available right now!")