pytorch 计算Parameter和FLOP的操作
这篇文章主要介绍了pytorch 计算Parameter和FLOP的操作,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
深度学习中,模型训练完后,查看模型的参数量和浮点计算量,在此记录下:
1 THOP
在pytorch中有现成的包thop用于计算参数数量和FLOP,首先安装thop:
1 | pip install thop |
注意安装thop时可能出现如下错误:
解决方法:
1 | pip install - - upgrade git + https: / / github.com / Lyken17 / pytorch - OpCounter.git # 下载源码安装 |
使用方法如下:
1 2 3 4 | from torchvision.models import resnet50 # 引入ResNet50模型 from thop import profile model = resnet50() flops, params = profile(model, input_size = ( 1 , 3 , 224 , 224 )) # profile(模型,输入数据) |
对于自己构建的函数也一样,例如shuffleNetV2
1 2 3 4 5 6 7 8 | from thop import profile from utils.ShuffleNetV2 import shufflenetv2 # 导入shufflenet2 模块 import torch model_shuffle = shufflenetv2(width_mult = 0.5 ) model = torch.nn.DataParallel(model_shuffle) # 调用shufflenet2 模型,该模型为自己定义的 flop, para = profile(model, input_size = ( 1 , 3 , 224 , 224 ),) print ( "%.2fM" % (flop / 1e6 ), "%.2fM" % (para / 1e6 )) |
更多细节,可参考thop GitHub链接: https://github.com/Lyken17/pytorch-OpCounter
2 计算参数
pytorch本身带有计算参数的方法
1 2 3 4 5 6 7 8 | from thop import profile from utils.ShuffleNetV2 import shufflenetv2 # 导入shufflenet2 模块 import torch model_shuffle = shufflenetv2(width_mult = 0.5 ) model = torch.nn.DataParallel(model_shuffle) total = sum ([param.nelement() for param in model.parameters()]) print ( "Number of parameter: %.2fM" % (total / 1e6 )) |
补充:pytorch: 计算网络模型的计算量(FLOPs)和参数量(Params)
计算量:
FLOPs,FLOP时指浮点运算次数,s是指秒,即每秒浮点运算次数的意思,考量一个网络模型的计算量的标准。
参数量:
Params,是指网络模型中需要训练的参数总数。
第一步:安装模块(thop)
1 | pip install thop |
第二步:计算
1 2 3 4 5 6 | import torch from thop import profile net = Model() # 定义好的网络模型 input = torch.randn( 1 , 3 , 112 , 112 ) flops, params = profile(net, (inputs,)) print ( 'flops: ' , flops, 'params: ' , params) |
注意:
输入input的第一维度是批量(batch size),批量的大小不回影响参数量, 计算量是batch_size=1的倍数
profile(net, (inputs,))的 (inputs,)中必须加上逗号,否者会报错
以上为个人经验,希望能给大家一个参考,