import torch
# torch版本
print('Pytorch version\t:', torch.__version__)
# cuda 版本
print('CUDA version\t:', torch.version.cuda)
# 使用的GPU数量及名称
for i in range(torch.cuda.device_count()):
print(f'GPU{i}\t\t:',torch.cuda.get_device_name(i))
# 当前使用的GPU
print('Current_device\t:', torch.cuda.current_device())
# GPU算力
print('Capability\t:', torch.cuda.get_device_capability(0))
# GPU是否可用
print('Is available\t:', torch.cuda.is_available())
# 设置device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Set device\t:', device)
# 获取设备属性
print('Properties\t:', torch.cuda.get_device_properties(0))
# 命令行查看设备
# nvidia-smi
打印结果:
Pytorch version : 1.5.0
CUDA version : 10.2
GPU0 : Tesla P100-PCIE-16GB
Current_device : 0
Capability : (6, 0)
Is available : True
Device : cuda:0
Properties : _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', major=6, minor=0, total_memory=16280MB, multi_processor_count=56)