如何计算caffe模型的参数量params与flops
程序员文章站
2022-05-28 23:48:16
...
最近需要比较不同模型的参数量,pytorch可以用一行代码解决,但是Caffe比较麻烦,做个记录~
# Pytorch
count = 0
for p in net.parameters():
count += p.data.nelement()
-
方法1. 通用型
文件 calc_params.py
import sys
sys.path.insert(0, "/home/ubuntu/workspace/caffe-advance/python")
import caffe
caffe.set_mode_cpu()
import numpy as np
from numpy import prod, sum
from pprint import pprint
def print_net_parameters_flops (deploy_file):
print ("Net: " + deploy_file)
net = caffe.Net(deploy_file, caffe.TEST)
flops = 0
typenames = ['Convolution', 'DepthwiseConvolution', 'InnerProduct']
print ("Layer-wise parameters: ")
print ('layer name'.ljust(20), 'Filter Shape'.ljust(20), \
'Output Size'.ljust(20), 'Layer Type'.ljust(20), 'Flops'.ljust(20))
for layer_name, blob in net.blobs.items():
if layer_name not in net.layer_dict:
continue
if net.layer_dict[layer_name].type in typenames:
cur_flops = 0.0
if net.layer_dict[layer_name].type in typenames[:2]:
cur_flops = (np.product(net.params[layer_name][0].data.shape) * \
blob.data.shape[-1] * blob.data.shape[-2])
else:
cur_flops = np.product(net.params[layer_name][0].data.shape)
print(layer_name.ljust(20),
str(net.params[layer_name][0].data.shape).ljust(20),
str(blob.data.shape).ljust(20),
net.layer_dict[layer_name].type.ljust(20), str(cur_flops).ljust(20))
# InnerProduct
if len(blob.data.shape) == 2:
flops += prod(net.params[layer_name][0].data.shape)
else:
flops += prod(net.params[layer_name][0].data.shape) * blob.data.shape[2] * blob.data.shape[3]
print ('layers num: ' + str(len(net.params.items())))
print ("Total number of parameters: " + str(sum([prod(v[0].data.shape) for k, v in net.params.items()])))
print ("Total number of flops: " + str(flops))
if __name__ == '__main__':
if len(sys.argv) != 2:
print ('Usage:')
print ('python calc_params.py deploy.prototxt')
exit()
deploy_file = sys.argv[1]
print_net_parameters_flops(deploy_file)
使用方法:只需修改第二行中的caffe的root路径为你自己的就可以了(可考虑加入自定义Python layer层)
python calc_params.py deploy.prototxt
方法2
这种方法可能会对某些层不识别,需要在不识别的层地方,跳过。
import sys
sys.path.insert(0,"python")
import caffe
model="models/bvlc_alexnet/deploy.prototxt"
def main():
net=caffe.Net(model,caffe.TEST)
params=0
flops=0
blobs=net.blobs
print("name param flops")
for item in net.params.items():
name,layer=item
c1=layer[0].count
c2=layer[1].count
b=blobs[name]
param=c1+c2
flop=param*b.width*b.height
print(name+" "+str(param)+" "+str(flop))
params+=param
flops+=flop
print("total params",params)
print("FLOPs:",flops)
if __name__ == '__main__':
main()