python----------->>>>>>>>>>>>统计caffe模型的模型参数量

#统计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))

  

posted @ 2021-03-17 11:06  水木清扬  阅读(210)  评论(0编辑  收藏  举报