如何计算caffe模型的参数量params与flops

最近需要比较不同模型的参数量，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()