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;
}

在最新版本的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()

它返回一个元组，其中第一个元素是空闲内存使用情况，第二个元素是总可用内存。

这对我很有用！

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