使用 pytorch 获取可用 GPU 内存的总量

Question

我正在使用 google colab 免费 Gpu 进行实验，并想知道有多少 GPU 内存可供使用，torch.cuda.memory_allocated() 返回当前占用的 GPU 内存，但我们如何使用 PyTorch 确定总可用内存。

Answer 1

PyTorch 可以为您提供总的、保留的和分配的信息：t = torch.cuda.get_device_properties(0).total_memoryr = torch.cuda.memory_reserved(0)a = torch.cuda.memory_allocated(0)f = r-a&nbsp; # free inside reserved与 NVIDIA 的 Python 绑定可以为您提供整个 GPU 的信息（在这种情况下，0 表示第一个 GPU 设备）：from pynvml import *nvmlInit()h = nvmlDeviceGetHandleByIndex(0)info = nvmlDeviceGetMemoryInfo(h)print(f'total&nbsp; &nbsp; : {info.total}')print(f'free&nbsp; &nbsp; &nbsp;: {info.free}')print(f'used&nbsp; &nbsp; &nbsp;: {info.used}')点安装pynvml您可以检查nvidia-smi以获取内存信息。您可以使用nvtop，但需要从源代码安装此工具（在撰写本文时）。另一个可以检查内存的工具是 gpustat ( pip3 install gpustat)。如果您想使用 C++ cuda：include <iostream>#include "cuda.h"#include "cuda_runtime_api.h"&nbsp;&nbsp;using namespace std;&nbsp;&nbsp;int main( void ) {&nbsp; &nbsp; int num_gpus;&nbsp; &nbsp; size_t free, total;&nbsp; &nbsp; cudaGetDeviceCount( &num_gpus );&nbsp; &nbsp; for ( int gpu_id = 0; gpu_id < num_gpus; gpu_id++ ) {&nbsp; &nbsp; &nbsp; &nbsp; cudaSetDevice( gpu_id );&nbsp; &nbsp; &nbsp; &nbsp; int id;&nbsp; &nbsp; &nbsp; &nbsp; cudaGetDevice( &id );&nbsp; &nbsp; &nbsp; &nbsp; cudaMemGetInfo( &free, &total );&nbsp; &nbsp; &nbsp; &nbsp; cout << "GPU " << id << " memory: free=" << free << ", total=" << total << endl;&nbsp; &nbsp; }&nbsp; &nbsp; return 0;}

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

Answer 3

这对我有用！def get_memory_free_MiB(gpu_index):&nbsp; &nbsp; pynvml.nvmlInit()&nbsp; &nbsp; handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index))&nbsp; &nbsp; mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)&nbsp; &nbsp; return mem_info.free // 1024 ** 2