Python get gpu memory

def get_gpu_memory_map() -> Dict[str, int]:
    """Get the current gpu usage.

    Return:
        A dictionary in which the keys are device ids as integers and
        values are memory usage as integers in MB.
    """
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader",],
        encoding="utf-8",
        # capture_output=True,          # valid for python version >=3.7
        stdout=PIPE,
        stderr=PIPE,  # for backward compatibility with python version 3.6
        check=True,
    )
    # Convert lines into a dictionary
    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]
    gpu_memory_map = {f"gpu_{index}": memory for index, memory in enumerate(gpu_memory)}
    return gpu_memory_map 

def get_gpu_memory_map():
    """Get the current gpu usage.

    Returns
    -------
    usage: dict
        Keys are device ids as integers.
        Values are memory usage as integers in MB.
    """
    result = subprocess.check_output(
        [
            'nvidia-smi', '--query-gpu=memory.used',
            '--format=csv,nounits,noheader'
        ])

    return int(result) 

def get_gpu_memory_usage(ctx: Union[mx.context.Context, List[mx.context.Context]]) -> Dict[int, Tuple[int, int]]:
    """
    Returns used and total memory for GPUs identified by the given context list.

    :param ctx: List of MXNet context devices.
    :return: Dictionary of device id mapping to a tuple of (memory used, memory total).
    """
    if isinstance(ctx, mx.context.Context):
        ctx = [ctx]
    ctx = [c for c in ctx if c.device_type == 'gpu']
    if not ctx:
        return {}

    memory_data = {}  # type: Dict[int, Tuple[int, int]]
    for c in ctx:
        try:
            free, total = mx.context.gpu_memory_info(device_id=c.device_id)  # in bytes
            used = total - free
            memory_data[c.device_id] = (used * 1e-06, total * 1e-06)
        except mx.MXNetError:
            logger.exception("Failed retrieving memory data for gpu%d", c.device_id)
            continue
    log_gpu_memory_usage(memory_data)
    return memory_data 

def get_gpu_memory_map():
    result = subprocess.check_output(
        [
            'nvidia-smi', '--query-gpu=memory.used',
            '--format=csv,nounits,noheader'
        ], encoding='utf-8')
    gpu_memory = [int(x) for x in result.strip().split('\n')]
    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))
    return gpu_memory_map 

def get_gpu_memory(pid=os.getpid()):
    info = os.popen('nvidia-smi').read()
    for line in info.split('\n'):
        t = line.strip().split()
        try:
            cur_pid = int(t[2])
            if cur_pid != pid:
                continue
            gpu_mem = int(t[5][:-3])
            return gpu_mem
        except:
            continue

    return 0 

def get_gpu_memory(gpus):
    """
    Small heuristic to calculate the amount of memory needed for each GPU in case of multi-gpu training.
    """
    if len(gpus) >= 5:
        return 0.4
    elif len(gpus) >= 3:
        return 0.5
    elif len(gpus) == 2:
        return 0.6
    else:
        return 0.6 

def get_gpu_memory():
  torch.cuda.synchronize()
  opts = [
      'nvidia-smi', '-q', '--gpu=' + str(0), '|', 'grep', '"Used GPU Memory"'
  ]
  cmd = str.join(' ', opts)
  ps = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.STDOUT)
  output = ps.communicate()[0].decode('utf-8')
  output = output.split("\n")[1].split(":")
  consumed_mem = int(output[1].strip().split(" ")[0])
  return consumed_mem 

def get_gpu_memory_map():
        result = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
            encoding="utf-8",
        )
        gpu_memory = [int(x) for x in result.strip().split("\n")]
        gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))
        return gpu_memory_map 

def get_gpu_memory_usage():
    """
    Grab nvidia-smi output and return a dictionary of the memory usage.
    """
    
    data = {'devCount':0, 'memTotal':0, 'memUsed':0, 'memFree':0, 'pwrDraw':0.0, 'pwrLimit':0.0, 'load':0.0}
    
    q_flag   = '--query-gpu=memory.used,memory.total,memory.free,power.draw,power.limit,utilization.gpu'
    fmt_flag = '--format=csv,noheader,nounits'
    try:
        p = subprocess.Popen(['nvidia-smi', q_flag, fmt_flag], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output, err = p.communicate()
    except (OSError, ValueError) as e:
        pass
    else:
        # Parse the ouptut and turn everything into something useful, if possible
        lines = output.split('\n')[:-1]
        for line in lines:
            used, total, free, draw, limit, load = line.split(',')
            data['devCount'] += 1
            data['memTotal'] += int(total, 10)*1024
            data['memUsed']  += int(used, 10)*1024
            data['memFree']  += int(free, 10)*1024
            try:
                data['pwrDraw'] += float(draw)
                data['pwrLimit'] += float(limit)
            except ValueError:
                pass
            try:
                data['load'] += float(load)
            except ValueError:
                pass
        # Convert the load to an average
        data['load'] /= data['devCount']
    return data 

def get_gpu_memory():
    torch.cuda.synchronize()
    opts = [
        'nvidia-smi', '-q', '--gpu=' + str(1), '|', 'grep', '"Used GPU Memory"'
    ]
    cmd = str.join(' ', opts)
    ps = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    output = ps.communicate()[0].decode('utf-8')
    output = output.split("\n")[0].split(":")
    consumed_mem = int(output[1].strip().split(" ")[0])
    return consumed_mem 

def get_gpu_memory_usage():
    if sys.platform == 'win32':
        # TODO: add windows support
        return None
    cuda_home = find_cuda()
    if cuda_home is None:
        return None
    cuda_home = Path(cuda_home)
    source = """
    #include <cuda_runtime.h>
    #include <iostream>
    int main(){
        int nDevices;
        cudaGetDeviceCount(&nDevices);
        size_t free_m, total_m;
        // output json format.
        std::cout << "[";
        for (int i = 0; i < nDevices; i++) {
            cudaSetDevice(i);
            cudaMemGetInfo(&free_m, &total_m);
            std::cout << "[" << free_m << "," << total_m << "]";
            if (i != nDevices - 1)
                std::cout << "," << std::endl;
        }
        std::cout << "]" << std::endl;
        return 0;
    }
    """
    with tempfile.NamedTemporaryFile('w', suffix='.cc') as f:
        f_path = Path(f.name)
        f.write(source)
        f.flush()
        try:
            # TODO: add windows support
            cmd = (
                f"g++ {f.name} -o {f_path.stem} -std=c++11"
                f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart")
            print(cmd)
            subprocess.check_output(cmd, shell=True, cwd=f_path.parent)
            cmd = f"./{f_path.stem}"
            usages = subprocess.check_output(
                cmd, shell=True, cwd=f_path.parent).decode()
            usages = json.loads(usages)
            return usages
        except:
            return None
    return None 

def get_gpu_memory_usage():
    if sys.platform == "win32":
        # TODO: add windows support
        return None
    cuda_home = find_cuda()
    if cuda_home is None:
        return None
    cuda_home = Path(cuda_home)
    source = """
    #include <cuda_runtime.h>
    #include <iostream>
    int main(){
        int nDevices;
        cudaGetDeviceCount(&nDevices);
        size_t free_m, total_m;
        // output json format.
        std::cout << "[";
        for (int i = 0; i < nDevices; i++) {
            cudaSetDevice(i);
            cudaMemGetInfo(&free_m, &total_m);
            std::cout << "[" << free_m << "," << total_m << "]";
            if (i != nDevices - 1)
                std::cout << "," << std::endl;
        }
        std::cout << "]" << std::endl;
        return 0;
    }
    """
    with tempfile.NamedTemporaryFile("w", suffix=".cc") as f:
        f_path = Path(f.name)
        f.write(source)
        f.flush()
        try:
            # TODO: add windows support
            cmd = (
                f"g++ {f.name} -o {f_path.stem} -std=c++11"
                f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart"
            )
            print(cmd)
            subprocess.check_output(cmd, shell=True, cwd=f_path.parent)
            cmd = f"./{f_path.stem}"
            usages = subprocess.check_output(
                cmd, shell=True, cwd=f_path.parent
            ).decode()
            usages = json.loads(usages)
            return usages
        except Exception:
            return None
    return None