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()-----------')

模型参数量统计 

for name, param in model.state_dict().items():
    print('name:{} param.size():{} param.numel():{} param.dtype:{}'.format(name, param.size(), param.numel(), param.dtype))

• 其中param.size()返回的是tensor的shape，而param.numel()返回的是size的乘积
• jit_model.state_dict()返回的是一个有序字典OrderedDict, 包含了所有可学习参数和不可学习参数；model.parameters()返回的是一个生成器，该生成器中只保存了可学习、可被优化器更新的参数的具体的参数；model.named_parameters()返回的是一个生成器(generator)，该生成器中只保存了可学习、可被优化器更新的参数的参数名和具体的参数
• 不可学习参数将会通过Module.register_parameter()注册在self._buffers中，self._buffers是一个OrderedDict，例如bn1层中的参数running_mean、running_var和num_batches_tracked均是不可学习参数
• 可学习参数也可叫做模型参数，其就是要参与学习和更新的，特别注意这里的参数更新是指在优化器的optim.step步骤里更新参数，即需要反向传播更新的参数，使用nn.parameter.Parameter()创建的变量是可学习参数(模型参数)

快速把GPU跑满：

 
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()))
 
dim = 10000
device_cnt = 2
a = torch.randn(dim, dim)
b = torch.randn(dim, dim)

gpu_a = [a.to(torch.device("cuda:{}".format(i) if torch.cuda.is_available() else "cpu")) for i in range(device_cnt)]
gpu_b = [b.to(torch.device("cuda:{}".format(i) if torch.cuda.is_available() else "cpu")) for i in range(device_cnt)]

while True:
    for i in range(device_cnt):
        c = gpu_a[i] @ gpu_b[i]

部分转载自Pytorch中state_dict()、named_parameters()和parameters()的区别_Jiyang@UESTC的博客-CSDN博客