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

约定不等于承诺〃 2024-04-18 07:59 60阅读 0赞

## 一、脚本 ##

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 calc_params.py  deploy.prototxt