PyTorch获取GPU信息(设备id、名称、显存)
在PyTorch中,我们可以使用如下代码获取GPU信息:
import torch
def gpu_info() -> str:
info = ''
for id in range(torch.cuda.device_count()):
p = torch.cuda.get_device_properties(id)
info += f'CUDA:{id} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n'
return info[:-1]
if __name__ == '__main__':
print(gpu_info())
本文版权,除注明引用的部分外,归作者所有。本文严禁商业用途的转载。非商业用途的转载需在网页明显处署上作者名称及原文链接。