torch.cuda 相关函数

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)