gpu_select.py

import os
import pandas as pd
from io import StringIO

def query_gpu(qargs=[]):
    '''
    qargs:
        query arguments
    return:
        a list of dict
    Querying GPUs infos
    查询GPU信息
    '''
    qargs = ['index', 'gpu_name', 'memory.free', 'memory.total', 'power.draw', 'power.limit',"utilization.gpu"] + qargs
    cmd = 'nvidia-smi --query-gpu={} --format=csv,noheader'.format(','.join(qargs))
    results = os.popen(cmd).read()
    csv_str = ",".join(qargs)+"\n"+results
    csv_str = csv_str.replace(" MiB", "")
    csv_str = csv_str.replace(" W", "")
    return pd.read_csv(StringIO(csv_str))
    # return csv_str
    
GPU_list=query_gpu()
# GPU_list["memory.free"].idxmax()
GPU_list
device = torch.device("cuda:{}".format(GPU_list["memory.free"].idxmax()))