模型计算量FLOPs的计算

读了一下之前的论文，对于计算量主要有Madds和MFlops两个概念。shufflenet的论文用的是Flops，Mobilenet用的是Madds，Flops应该是Madds的两倍，具体可参考

https://blog.csdn.net/shwan_ma/article/details/84924142

https://www.zhihu.com/question/65305385/answer/451060549

计算公式为：

 

下面实现了tf计算FLOPs的脚本代码，主要考虑了conv和depthwise，其他结构没有用到。输入为带weight的pb文件：

from tensorflow.python.framework import tensor_util
from google.protobuf import text_format
import tensorflow as tf
from tensorflow.python.platform import gfile
from tensorflow.python.framework import tensor_util
from numpy import prod, sum
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

flops = 0
GRAPH_PB_PATH = './net_model.pb' #path to your .pb file
with tf.Session() as sess:
	print("load graph")
	with gfile.FastGFile(GRAPH_PB_PATH,'rb') as f:
		graph_def = tf.GraphDef()
		graph = tf.get_default_graph()
		graph_def.ParseFromString(f.read())
				
		for node in graph_def.node :
		    if node.op == 'RefSwitch':
			    node.op = 'Switch'
			    for index in range(len(node.input)):
				    if 'moving_' in node.input[index]:
				    	node.input[index] = node.input[index] + '/read'
		    elif node.op == 'AssignSub':
			    node.op = 'Sub'
			    if 'use_locking' in node.attr: 
				    del node.attr['use_locking']
		
		tf.import_graph_def(graph_def, name='')
		num_layers = 0
		for op in graph.get_operations():
			if(op.type == "Conv2D" or op.type == "DepthwiseConv2dNative"):
				flops += op.outputs[0].shape[1] * op.outputs[0].shape[2] * prod(op.inputs[1].shape)
				#flops += op.outputs[0].shape[1] * op.outputs[0].shape[2] * prod(op.inputs[1].shape) + prod(prod(op.outputs[0].shape))
				print(op.name, op.outputs[0].shape, op.inputs[1].shape)
				num_layers += 1
		print("Total layers: ", num_layers)
		print("FLOPs: ", flops)