日本免费高清视频-国产福利视频导航-黄色在线播放国产-天天操天天操天天操天天操|www.shdianci.com

學(xué)無先后，達者為師

網(wǎng)站首頁 編程語言 正文

如何計算?tensorflow?和?pytorch?模型的浮點運算數(shù)_python

作者：浩哥依然 ? 更新時間： 2022-12-24 編程語言

本文主要討論如何計算 tensorflow 和 pytorch 模型的 FLOPs。如有表述不當(dāng)之處歡迎批評指正。歡迎任何形式的轉(zhuǎn)載，但請務(wù)必注明出處。

1. 引言

FLOPs 是 floating point operations 的縮寫，指浮點運算數(shù)，可以用來衡量模型/算法的計算復(fù)雜度。本文主要討論如何在 tensorflow 1.x, tensorflow 2.x 以及 pytorch 中利用相關(guān)工具計算對應(yīng)模型的 FLOPs。

2. 模型結(jié)構(gòu)

為了說明方便，先搭建一個簡單的神經(jīng)網(wǎng)絡(luò)模型，其模型結(jié)構(gòu)以及主要參數(shù)如表1 所示。

表 1 模型結(jié)構(gòu)及主要參數(shù)

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

用 tensorflow（實際使用 tensorflow 中的 keras 模塊）實現(xiàn)該模型的代碼為：

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 實現(xiàn)該模型的代碼為：

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

本節(jié)討論的版本具體為：tensorflow 1.12.0, tensorflow 2.3.1 以及 pytorch 1.10.1+cu102。

3.1. tensorflow 1.12.0

在 tensorflow 1.12.0 環(huán)境中，可以使用以下代碼計算模型的 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 環(huán)境中，可以使用以下代碼計算模型的 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 環(huán)境中，可以使用以下代碼計算模型的 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. 結(jié)果對比

三者計算出的 FLOPs 分別為：
tensorflow 1.12.0：

tensorflow 2.3.1：

pytorch 1.10.1：


可以看到 tensorflow 1.12.0 和 tensorflow 2.3.1 的結(jié)果基本在同一個量級，而與 pytorch 1.10.1 計算出來的相差甚遠。但如果將上述模型結(jié)構(gòu)改為只包含第一層 Conv2D，三者計算出來的 FLOPs 卻又是一致的。所以推斷差異主要來自于 GRU 的 FLOPs。如讀者知道其中詳情，還請不吝賜教。

4. 總結(jié)

本文給出了在 tensorflow 1.x, tensorflow 2.x 以及 pytorch 中利用相關(guān)工具計算模型 FLOPs 的方法，但從本文所使用的測試模型來看， tensorflow 與 pytorch 統(tǒng)計出的結(jié)果相差甚遠。當(dāng)然，也可以根據(jù)網(wǎng)絡(luò)層的類型及其對應(yīng)的參數(shù)，推導(dǎo)計算出每個網(wǎng)絡(luò)層所需的 FLOPs。

原文鏈接：https://blog.csdn.net/wjrenxinlei/article/details/127973081

欄目分類
最近更新