找到内存足够的GPU

我想通过编程找出可用的GPU及其当前内存使用情况，并根据其内存可用性使用其中一个GPU。我想在PyTorch做这件事

我在这方面看到了以下解决方案：

导入torch.cuda作为cutorch
对于范围内的i（cutorch.device_count（））：
如果cutorch.getMemoryUsage（i）>MEM：
opts.gpuID=i
打破

---

但它在PyTorch 0.3.1中不起作用（没有调用函数，

getMemoryUsage

）。我对基于PyTorch（使用库函数）的解决方案感兴趣。任何帮助都将不胜感激。

在您提供的网页中，有一个答案：

#/usr/bin/env python
#编码：utf-8
导入子流程
def get_gpu_memory_map（）：
“”“获取当前gpu使用情况。
退换商品
-------
用法：dict
键是作为整数的设备ID。
值是以MB为单位的整数形式的内存使用量。
"""
结果=子流程。检查\u输出(
[
'nvidia smi'，'--查询gpu=memory.used'，
“--format=csv，nounits，noheader”
])
#将行转换成字典
gpu_内存=[result.strip（）.split（'\n'）中x的int（x）]
gpu内存映射=dict（zip（范围（len（gpu内存）），gpu内存））
返回gpu内存映射
打印获取gpu内存映射（）

改进代码以显示cutorch.getMemoryUsage（i）的结果，并将该结果也放入您的问题中。换句话说，编辑您的问题以准确解释“它不起作用”的含义谢谢！我已经看到了这个解决方案，但我正在寻找一个不运行

nvidiasmi

命令的解决方案。更感兴趣的是了解是否有任何基于Pythorch的解决方案。您需要在nvidia smi的

查询之后添加
result=result.decode（'utf-8'）
，才能在python3下工作。