caffe计算模型中每层参数的个数和FLOPs

 

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

输出

name param flops
conv1 34944 105705600
conv2 307456 224135424
conv3 885120 149585280
conv4 663936 112205184
conv5 442624 74803456
fc6 37752832 37752832
fc7 16781312 16781312
fc8 4097000 4097000
('total params', 60965224)
('FLOPs:', 725066088)

参考:

CNN中parameters和FLOPs计算