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)