找到具有足够内存的GPU

时间:2018-04-01 07:08:51

标签: gpu pytorch

我想以编程方式找出可用的GPU及其当前的内存使用情况,并根据其内存可用性使用其中一个GPU。我想在PyTorch中这样做。

我在此post中看到了以下解决方案:

import torch.cuda as cutorch

for i in range(cutorch.device_count()):
    if cutorch.getMemoryUsage(i) > MEM: 
        opts.gpuID = i
        break

但它在PyTorch 0.3.1中不起作用(没有名为getMemoryUsage的函数)。我对基于PyTorch(使用库函数)的解决方案感兴趣。任何帮助将不胜感激。

1 个答案:

答案 0 :(得分:2)

在您提供的网页中,有一个答案:

#!/usr/bin/env python
# encoding: utf-8

import subprocess

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'
        ])
    # Convert lines into a dictionary
    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

print get_gpu_memory_map()