我正在使用 google colab 免费 Gpu 进行实验,并想知道有多少 GPU 内存可供使用,torch.cuda.memory_allocated() 返回当前占用的 GPU 内存,但我们如何使用 PyTorch 确定总可用内存。
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;
}
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)