找到一个有足够内存的 GPU

【问题标题】：Find a GPU with enough memory找到一个有足够内存的 GPU
【发布时间】：2018-09-10 17:58:09
【问题描述】：

我想以编程方式找出可用的 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（使用库函数）的解决方案感兴趣。任何帮助将不胜感激。

【问题讨论】：

改进代码以显示cutorch.getMemoryUsage(i) 的结果并将结果也放入您的问题中。
换句话说，编辑您的问题以准确地解释“它不起作用”的含义

标签： gpu pytorch

【解决方案1】：

在你给出的网页中，有一个答案：

#!/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()

【讨论】：

谢谢！我已经看到了该解决方案，但我正在寻找一个不运行 nvidia-smi 命令的解决方案。更想知道是否有任何基于 PyTorch 的解决方案。
您需要在nvidia-smiquery 之后添加result = result.decode('utf-8') 才能在python3 下工作。