6种方法计算神经网络参数量Params、计算量FLOPs、Macs简单代码

方法1:统计模型参数量

total = sum([param.nelement() for param in model.parameters()])

print("Number of parameter: %.2fM" % (total/1e6))

方法2:统计flops和参数量

pip install thop

from thop import profile

dummy_input = torch.randn(1, 3, 32, 32)#.to(device)

flops, params = profile(model, (dummy_input,))

print('flops: ', flops, 'params: ', params)

print('flops: %.2f M, params: %.2f M' % (flops / 1000000.0, params / 1000000.0))

问题:当网络中有自定义参数时,flops和params就很有可能漏掉那部分参数

方法3:统计macs指标和参数量

pip install ptflops

from ptflops import get_model_complexity_info

macs, params = get_model_complexity_info(model, (3, 32, 32), as_strings=True, print_per_layer_stat=True)

print('{:<30}  {:<8}'.format('Computational complexity(macs): ', macs))

print('{:<30}  {:<8}'.format('Number of parameters(params): ', params))

问题:当网络中有自定义参数时,macs就很有可能漏掉那部分参数

方法4:统计模型参数所占内存

def get_parameter_number(net):

    total_num = sum(p.numel() for p in net.parameters())

    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)

    return {'Total': total_num/ 1000000.0, 'Trainable': trainable_num/ 1000000.0}

#查看网络参数

print(get_parameter_number(model))

方法5:stat(cpu统计)

pip install torchstat

from torchstat import stat

stat(model, (3, 32, 32))

问题:当网络中有自定义参数时,就很有可能漏掉那部分参数对应的统计量

方法6:summary网络结构对应参数(cuda上面统计)

pip install torchsummary

from torchsummary import summary

summary(model.cuda(),input_size=(3,32,32),batch_size=-1)

问题:当网络中有自定义参数时,就很有可能漏掉那部分参数

你可能感兴趣的:(神经网络,python,深度学习,pytorch)