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)+