pytorch和tensorflow关于神经网络的参数量估计和运算量估计（Parameters和FLOPs）

浮点运算数估计和参数量估计

简单地记录一下关于神经网络参数估计和计算量估计的心得。

Pytorch

Pytorch没有tensorflow那么繁琐，这里推荐使用Thop包来进行计算。下面是一个例子：

from thop import profile
from importlib import import_module

MODEL = import_module(args.model)
model = MODEL.Net(args).to('cpu') # 看情况需变更，取得可使用的model
input = torch.randn(1, 3, 1024) # 点云的输入，如果是图片则需要自行更改尺寸
flops, params = profile(model, inputs=(input, ))
print('flops: %.2f M, params: %.2f M' % (flops / 1000000.0, params / 1000000.0))

另一种是使用torch自带的计算功能进行计算，建议用统一的thop来计算。其他可参考链接：
知乎:Pytorch与Tensorflow计算网络参数量Params和运行时间

Tensorflow

tensorflow相对而言用的比较少，个人感觉比较繁琐。首先要理解tensorflow中graph和session代表什么，graph是纯粹的网络结构，session是输入的对应参数，感觉在torch中没有对应session的函数，可能和forward差不多。
tensorflow计算网上很多都是需要ckpt训练好的权重文件，这里给一个例子：

import tensorflow as tf
def stats_graph(graph):
	flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
	params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
	print('FLOPs: {} M;    Trainable params: {} M'.format(flops.total_float_ops / 1000000.0, params.total_parameters / 1000000.0))

if __name__=='__main__':
    with tf.Graph().as_default() as graph:
        inputs = tf.zeros((1,1024,3)) # 点云格式
        output, _ = get_model(inputs, tf.constant(True))
        stats_graph(graph)

总结

要先找到待计算神经网络的模型文件，新建main函数，导入对应的包，设置单独的输入，得到计算的结果。tensorflow问题还是比较多的，很多用1.x版本的网络模型如果你装了2.x的tensorflow就很难用，换包用很麻烦，还是直接创建新的对应python版本的虚拟环境，装一个相应版本的tensorflow-gpu进行编译和跑代码最快。用错误的版本编译得到的so等文件都要删除重新用对的版本编译才能得到正确的执行文件，不然更换了版本重新编译一般不会跑通。