tensorflow 模型浮点数计算量和参数量估计

写在前面的话

之前在pytorch统计参数量和计算量用stat非常的方便，现在转到tensorflow后发现tensorflow的参数量和计算量的统计相对来说就没那么方便了。
在网上也搜了一些相关教程，发现用profiler来进行统计是最方便的。

本文参考了：tensorflow 模型浮点数计算量和参数量估计和TensorFlow程序分析（profile）实战，他们都整理的都非常好，我在这里重新整理一遍方便自己看。

计算量、参数量实验

stats_graph.py

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: {};    Trainable params: {}'.format(flops.total_float_ops, params.total_parameters))

搭建一个Graph实验一下：

with tf.Graph().as_default()