#统计caffe算模型的参数量
from numpy import prod, sum
flops = 0
typenames = ['Convolution', 'BatchNorm']
for layer_name, blob in net.blobs.iteritems():
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))