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

问题描述

我正在使用Google Colab Free 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

与NVIDIA的Python绑定可以为您带来整个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 need 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屋！