caffe 获取模型flops

import sys
sys.path.insert(0,"python")
import caffe
 

def get_model_param_flops(model):
    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 = 0
        flop = 0
        if len(layer) == 2:
            c2=layer[1].count
            b=blobs[name]
            
            flop=(c1+c2)*b.width*b.height

        param=c1+c2
        print(name+" "+str(param)+" "+str(flop))
        params+=param
        flops+=flop
    print("total params",params)
    print("FLOPs:",flops)
	
	
if __name__ == '__main__':
    model="o_net_half2.prototxt"
    get_model_param_flops(model)