如何计算 tensorflow 和 pytorch 模型的浮点运算数

更新时间：2022年11月26日 16:57:10 作者：浩哥依然

FLOPs 是 floating point operations 的缩写，指浮点运算数，可以用来衡量模型/算法的计算复杂度。本文主要讨论如何在 tensorflow 1.x, tensorflow 2.x 以及 pytorch 中利用相关工具计算对应模型的 FLOPs,需要的朋友可以参考下

1. 引言

FLOPs 是 floating point operations 的缩写，指浮点运算数，可以用来衡量模型/算法的计算复杂度。本文主要讨论如何在 tensorflow 1.x, tensorflow 2.x 以及 pytorch 中利用相关工具计算对应模型的 FLOPs。

2. 模型结构

为了说明方便，先搭建一个简单的神经网络模型，其模型结构以及主要参数如表1 所示。

表 1 模型结构及主要参数

Layers	channels	Kernels	Strides	Units	Activation
Conv2D	32	(4,4)	(1,2)	\	relu
GRU	\	\	\	96	\
Dense	\	\	\	256	sigmoid

用 tensorflow（实际使用 tensorflow 中的 keras 模块）实现该模型的代码为：

from tensorflow.keras.layers import *
from tensorflow.keras.models import load_model, Model

def test_model_tf(Input_shape):
    # shape: [B, C, T, F]
    main_input = Input(batch_shape=Input_shape, name='main_inputs')
    
    conv = Conv2D(32, kernel_size=(4, 4), strides=(1, 2), activation='relu', data_format='channels_first', name='conv')(main_input)
    
    # shape: [B, T, FC]
    gru = Reshape((conv.shape[2], conv.shape[1] * conv.shape[3]))(conv)
    gru = GRU(units=96, reset_after=True, return_sequences=True, name='gru')(gru)
    
    output = Dense(256, activation='sigmoid', name='output')(gru)
    
    model = Model(inputs=[main_input], outputs=[output])
    
    return model

用 pytorch 实现该模型的代码为：

import torch
import torch.nn as nn

class test_model_torch(nn.Module):
    def __init__(self):
        super(test_model_torch, self).__init__()

        self.conv2d = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(4,4), stride=(1,2))
        self.relu = nn.ReLU()

        self.gru = nn.GRU(input_size=4064, hidden_size=96)

        self.fc = nn.Linear(96, 256)
        self.sigmoid = nn.Sigmoid()

    def forward(self, inputs):
        # shape: [B, C, T, F]
        out = self.conv2d(inputs)
        out = self.relu(out)
        
        # shape: [B, T, FC]
        batch, channel, frame, freq = out.size()
        out = torch.reshape(out, (batch, frame, freq*channel))
        out, _ = self.gru(out)
        
        out = self.fc(out)
        out = self.sigmoid(out)

        return out

3. 计算模型的 FLOPs

本节讨论的版本具体为：tensorflow 1.12.0, tensorflow 2.3.1 以及 pytorch 1.10.1+cu102。

3.1. tensorflow 1.12.0

在 tensorflow 1.12.0 环境中，可以使用以下代码计算模型的 FLOPs：

import tensorflow as tf
import tensorflow.keras.backend as K

def get_flops(model):
    run_meta = tf.RunMetadata()
    opts = tf.profiler.ProfileOptionBuilder.float_operation()

    flops = tf.profiler.profile(graph=K.get_session().graph,
                                run_meta=run_meta, cmd='op', options=opts)
 
    return flops.total_float_ops

if __name__ == "__main__":
    x = K.random_normal(shape=(1, 1, 100, 256))
    model = test_model_tf(x.shape)
    print('FLOPs of tensorflow 1.12.0:', get_flops(model))

3.2. tensorflow 2.3.1

在 tensorflow 2.3.1 环境中，可以使用以下代码计算模型的 FLOPs ：

import tensorflow.compat.v1 as tf
import tensorflow.compat.v1.keras.backend as K
tf.disable_eager_execution()

def get_flops(model):
    run_meta = tf.RunMetadata()
    opts = tf.profiler.ProfileOptionBuilder.float_operation()

    flops = tf.profiler.profile(graph=K.get_session().graph,
                                run_meta=run_meta, cmd='op', options=opts)
 
    return flops.total_float_ops

if __name__ == "__main__":
    x = K.random_normal(shape=(1, 1, 100, 256))
    model = test_model_tf(x.shape)
    print('FLOPs of tensorflow 2.3.1:', get_flops(model))

3.3. pytorch 1.10.1+cu102

在 pytorch 1.10.1+cu102 环境中，可以使用以下代码计算模型的 FLOPs（需要安装 thop）：

import thop

x = torch.randn(1, 1, 100, 256)
model = test_model_torch()
flops, _ = thop.profile(model, inputs=(x,))
print('FLOPs of pytorch 1.10.1:', flops * 2)

需要注意的是，thop 返回的是 MACs (Multiply–Accumulate Operations)，其等于 2 2 2 倍的 FLOPs，所以上述代码有乘 2 2 2 操作。

3.4. 结果对比

三者计算出的 FLOPs 分别为：
tensorflow 1.12.0：

tensorflow 2.3.1：

pytorch 1.10.1：

可以看到 tensorflow 1.12.0 和 tensorflow 2.3.1 的结果基本在同一个量级，而与 pytorch 1.10.1 计算出来的相差甚远。但如果将上述模型结构改为只包含第一层 Conv2D，三者计算出来的 FLOPs 却又是一致的。所以推断差异主要来自于 GRU 的 FLOPs。如读者知道其中详情，还请不吝赐教。

4. 总结

本文给出了在 tensorflow 1.x, tensorflow 2.x 以及 pytorch 中利用相关工具计算模型 FLOPs 的方法，但从本文所使用的测试模型来看， tensorflow 与 pytorch 统计出的结果相差甚远。当然，也可以根据网络层的类型及其对应的参数，推导计算出每个网络层所需的 FLOPs。

到此这篇关于计算 tensorflow 和 pytorch 模型的浮点运算数的文章就介绍到这了,更多相关tensorflow 和 pytorch浮点运算数内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:

Python实现列表转Excel表格的第一列
这篇文章主要为大家详细介绍了如何将Python中的列表转换为Excel表格的第一列,并通过案例和代码展示具体的操作步骤,希望可以帮助大家快速掌握这一技能
2024-04-04
PyCharm在新窗口打开项目的方法
今天小编就为大家分享一篇PyCharm在新窗口打开项目的方法，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2019-01-01
Django利用Channels+websocket开发聊天室完整案例
Channels是Django团队研发的一个给Django提供websocket支持的框架,使用它我们可以轻松开发需要长链接的实时通讯应用,下面这篇文章主要给大家介绍了关于Django利用Channels+websocket开发聊天室的相关资料,需要的朋友可以参考下
2023-06-06
OpenCV半小时掌握基本操作之模板匹配
这篇文章主要介绍了OpenCV基本操作之模板匹配，本文给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值，需要的朋友可以参考下
2021-09-09
Python中设置变量访问权限的方法
这篇文章主要介绍了Python中设置变量访问权限的方法,是Python学习当中的重要知识点,需要的朋友可以参考下
2015-04-04
python 图片验证码代码分享
python 图片验证码代码分享，需要的朋友可以参考下
2012-07-07
Python3 全自动更新已安装的模块实现
这篇文章主要介绍了Python3 全自动更新已安装的模块实现，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2020-01-01
python3使用GUI统计代码量
这篇文章主要为大家详细介绍了python3使用GUI统计代码量，具有一定的参考价值，感兴趣的小伙伴们可以参考一下
2019-09-09
python将excel转换为csv的代码方法总结
在本篇文章里小编给大家分享了关于python如何将excel转换为csv的实例方法和代码内容，需要的朋友们学习下。
2019-07-07
Django中的CACHE_BACKEND参数和站点级Cache设置
这篇文章主要介绍了Django中的CACHE_BACKEND参数和站点级Cache设置,Python是最具人气的Python web框架,需要的朋友可以参考下
2015-07-07