Pytorch查看GPU是否可用、模型参数量numel统计

查看是否可用

import torch
import os

print('torch.cuda.is_available() = {}'.format(torch.cuda.is_available()))
print("torch.cuda.device_count()={}".format(torch.cuda.device_count()))
print("torch.cuda.get_device_name(0)={}".format(torch.cuda.get_device_name(0)))
print("torch.cuda.current_device()={}".format(torch.cuda.current_device()))
 
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
a_cpu = torch.Tensor(2, 3)
print('a_cpu.device={}'.format(a_cpu.device))
# device(type='cpu')
 
a_cuda = a_cpu.to(device)
# tensor([[-2.1800e-01,  4.5737e-41,  2.3351e-37],
#         [ 0.0000e+00,  4.4842e-44,  0.0000e+00]], device='cuda:0')
print('a_cuda.device={}'.format(a_cuda.device))

print('-----------os.environ.items()-----------')
for k,v in os.environ.items():
    print('os.environ.items_key={} v={}'.format(k,v))
print('-----------os.environ.items()-----------')