使用 pytorch 获取可用的 GPU 内存总量 [英] Get total amount of free GPU memory and available using pytorch

问题描述

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

I'm using google colab free Gpu's for experimentation and wanted to know how much GPU Memory available to play around, torch.cuda.memory_allocated() returns the current GPU memory occupied, but how do we determine total available memory using PyTorch.

推荐答案

PyTorch 可以为您提供全部、保留和分配的信息:

PyTorch can provide you total, reserved and allocated info:

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 设备):

Python bindings to NVIDIA can bring you the info for the whole GPU (0 in this case means first GPU device):

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).

You may check the nvidia-smi to get memory info. You may use nvtop but this tool needs to be installed from source (at the moment of writing this). Another tool where you can check memory is gpustat (pip3 install gpustat).

如果您想使用 C++ cuda:

If you would like to use 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 获取可用的 GPU 内存总量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！