Получить общий объем свободной памяти GPU и доступный с помощью Pytorch

Question

Я использую Google Colab Free Gpu для экспериментов и хотел узнать, сколько памяти GPU доступно для воспроизведения, torch.cuda.memory_allocated () возвращает занятую текущую память GPU, но как определить общий объем доступной памяти с помощью PyTorch? .

Answer 1

PyTorch может предоставить вам общую, кэшированную и выделенную информацию:

t = torch.cuda.get_device_properties(0).total_memory
c = torch.cuda.memory_cached(0)
a = torch.cuda.memory_allocated(0)
f = c-a  # free inside cache

Привязки Python к NVIDIA могут принести вам информацию для всего графического процессора (0 в данном случае означает первое устройство с графическим процессором):

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