使用PyTorch获取GPU的总空闲内存和可用内存

Question

我正在使用Google Colab提供的免费GPU进行实验，想知道可以用多少GPU内存进行测试。通过调用torch.cuda.memory_allocated()可以返回当前已经占用的GPU内存，但是如何确定PyTorch中总共可用的内存呢？

- Hari Prasad

Answer 1

PyTorch可以为您提供总信息、保留信息和分配信息：

t = torch.cuda.get_device_properties(0).total_memory
r = torch.cuda.memory_reserved(0)
a = torch.cuda.memory_allocated(0)
f = r-a  # free inside reserved

Python绑定到NVIDIA可以为您提供整个GPU的信息（在本例中，0表示第一个GPU设备）：

from pynvml import *
nvmlInit()
h = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(h)
print(f'total    : {info.total}')
print(f'free     : {info.free}')
print(f'used     : {info.used}')

_{pip install pynvml}

您可以使用nvidia-smi来获取内存信息。 您可以使用nvtop，但是此工具需要从源代码安装（在撰写本文时）。 另一个检查内存的工具是gpustat（pip3 install gpustat）。

如果您想使用C++ cuda：

include <iostream>
#include "cuda.h"
#include "cuda_runtime_api.h"
  
using namespace std;
  
int main( void ) {
    int num_gpus;
    size_t free, total;
    cudaGetDeviceCount( &num_gpus );
    for ( int gpu_id = 0; gpu_id < num_gpus; gpu_id++ ) {
        cudaSetDevice( gpu_id );
        int id;
        cudaGetDevice( &id );
        cudaMemGetInfo( &free, &total );
        cout << "GPU " << id << " memory: free=" << free << ", total=" << total << endl;
    }
    return 0;
}

- prosti

---

torch.cuda.memory_cached 已更名为 torch.cuda.memory_reserved。 - Kallzvx

已更新@Kallzvx。如果有什么不对的地方，请告诉我。 - prosti

注意：当内存由其他用户/进程分配时，total_memory + reserved/allocated 的计算结果可能不准确。 - krassowski

请使用import pynvml代替from pynvml import *，否则可能会与其他代码发生冲突。例如，modeling_roberta.py会抛出TypeError: '_ctypes.UnionType' object is not subscriptable。pynvml.nvmlInit()， h = pynvml.nvmlDeviceGetHandleByIndex(0)， info = pynvml.nvmlDeviceGetMemoryInfo(h)。 - user2585501

Answer 2

在最近的PyTorch版本中，您还可以使用torch.cuda.mem_get_info：

https://pytorch.org/docs/stable/generated/torch.cuda.mem_get_info.html#torch.cuda.mem_get_info

torch.cuda.mem_get_info()

它返回一个元组，其中第一个元素是可用的自由内存使用量，第二个元素是总可用内存。

- Iman

---

这个回答比被接受的回答更好（使用 total_memory+ 保留/分配），因为它在其他进程/用户共享GPU并占用内存时提供正确的数字。 - krassowski

在旧版本的PyTorch中存在缺陷，它会忽略设备参数并始终返回当前设备信息。解决方法是使用上下文管理器：with torch.cuda.device(device): info = torch.cuda.mem_get_info()参见：https://github.com/pytorch/pytorch/issues/76224 - אלימלך שרייבר

示例用法请 - Nathan B

@NathanB 添加了使用示例。 - Iman

Answer 3

这对我很有用！

def get_memory_free_MiB(gpu_index):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index))
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    return mem_info.free // 1024 ** 2

- Peter Pack