欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

如何计算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()