[人工智能]Caffe CNN模型FLOPs 计算脚本

简单估算caffe模型参数量和计算量的脚本，包含Convolution,?MaxPooling,?Eltwise和InnerProduct层

import os, sys
os.environ['GLOG_minloglevel'] = '2'

CAFFE_ROOT = "/home/hp/Documents/caffe-master"
if os.path.join(CAFFE_ROOT, 'python') not in sys.path:
    sys.path.insert(0, os.path.join(CAFFE_ROOT, 'python'))

import caffe
import google.protobuf.text_format
import google.protobuf as pb
from caffe.proto.caffe_pb2 import NetParameter

model="models/bvlc_alexnet/deploy.prototxt"
 
def main():
    net=caffe.Net(model,caffe.TEST)
    params_num=0
    flops=0
    net_info = NetParameter()
    with open(model, 'r') as fp:
        pb.text_format.Merge(fp.read(), net_info)

    layer_dict = {layer.name : layer for layer in net_info.layer}
    blobs=net.blobs
    params=net.params
    print("name".ljust(25)+"type".ljust(25)+"param".ljust(25)+"FLOPs".ljust(25))
    for item in layer_dict.items():
        name,layer=item
        type=layer.type
        if type == 'Convolution':
            b=blobs[name]
            weights=params[name][0]
            c1=weights.count/layer.convolution_param.group
            c2=0
            if len(params[name]) > 1:
                bias=params[name][1]
                c2=bias.count
            param=c1+c2
            flop=2*param*b.width*b.height*b.num+b.count #Convolution+ReLU
            print(name.ljust(25)+str(type).ljust(25)+str(param).ljust(25)+str(flop).ljust(25))
            params_num+=param
            flops+=flop
        elif type == 'Pooling': #MaxPooling
            b=blobs[name]
            k=layer.pooling_param.kernel_size
            param=0
            flop=k*k*b.count
            print(name.ljust(25)+str(type).ljust(25)+str(param).ljust(25)+str(flop).ljust(25))
            params_num+=param
            flops+=flop
        elif type == 'Eltwise':
            b=blobs[name]
            param=0
            flop=b.count
            print(name.ljust(25)+str(type).ljust(25)+str(param).ljust(25)+str(flop).ljust(25))
            params_num+=param
            flops+=flop
        elif type == 'InnerProduct':
            b=blobs[name]
            weights=params[name][0]
            c1=weights.count
            c2=0
            if len(params[name]) > 1:
                bias=params[name][1]
                c2=bias.count
            param=c1+c2
            flop=2*b.count*weights.channels
            print(name.ljust(25)+str(type).ljust(25)+str(param).ljust(25)+str(flop).ljust(25))
            params_num+=param
            flops+=flop
    print("total params(M)",params_num/1000/1000)
    print("GFLOPs:",flops/1000/1000/1000)
if __name__ == '__main__':
    main()

hp@hp-HP-Pavilion-Laptop-15-ck0xx:~/Documents/caffe-master$ conda activate caffe_env
(caffe_env) hp@hp-HP-Pavilion-Laptop-15-ck0xx:~/Documents/caffe-master$ python FLOPs.py
name                     type                     param                    FLOPs                    
fc6                      InnerProduct             37752832                 754974720                
fc7                      InnerProduct             16781312                 335544320                
pool5                    Pooling                  0                        829440                   
fc8                      InnerProduct             4097000                  81920000                 
conv3                    Convolution              885120.0                 2992354560.0             
conv5                    Convolution              221440.0                 748899840.0              
conv2                    Convolution              153856.0                 2245086720.0             
pool2                    Pooling                  0                        3893760                  
pool1                    Pooling                  0                        6298560                  
conv1                    Convolution              34944.0                  2117016000.0             
conv4                    Convolution              332160.0                 1123349760.0             
total params(M) 60.258663999999996
GFLOPs: 10.41016768
(caffe_env) hp@hp-HP-Pavilion-Laptop-15-ck0xx:~/Documents/caffe-master$ 

?如有错误欢迎指正

参考文献：https://arxiv.org/abs/1611.06440v2