神经网络中FLOPs和MACs的计算（基于thop和fvcore.nn）

最新推荐文章于 2024-05-22 21:38:52 发布

Chauncey_Wang

最新推荐文章于 2024-05-22 21:38:52 发布

阅读量4.6k

点赞数 5

分类专栏：笔记文章标签：神经网络深度学习 cnn

本文链接：https://blog.csdn.net/weixin_39518984/article/details/125081077

版权

笔记专栏收录该内容

7 篇文章 0 订阅

订阅专栏

以

输入为(1,1,200,3)的张量

卷积取

nn.Conv2d(1, 64, kernel_size=(8, 1), stride=(2, 1), padding=(0, 0))

为例。

先计算输出的形状

公式为 $\frac{n+2p-k}{s}+1$

H上为(200+0-8)/2+1=97

W上依然是3

所以输出的形状是(1,64,97,3)

卷积的本质是wx+b，

但是实际计算过程中，是直接w和x一一对应的乘起来，并且将结果都加起来

计算FLOPs时，一般会忽略b，而MACs并不会忽略b

所以对于一个卷积，对应的FLOPs为

97*3*(8*1)*64=148992

而对应的MACs为

97*3*(8*1)*64+97*3*64=167616

后一个97*3*64，就是对应b的数量

用代码计算的话，可以用thop计算MACs，fvcore.nn计算FLOPs

import torch
import torch.nn as nn
import torch.nn.functional as F
from thop import profile
from fvcore.nn import FlopCountAnalysis, parameter_count_table
class net(torch.nn.Module):
    def __init__(self, image_channels=1, n_classes=6):
        super(net, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(image_channels, 64, kernel_size=(8, 1), stride=(2, 1), padding=(0, 0)),
        )

    def forward(self, x):
        cnn_x = self.cnn(x)
        return cnn_x

model = net()
x = torch.randn(1, 1, 200, 3)

macs, params = profile(model, inputs=(x, ))  # ,verbose=False
print("MACs", macs)
print("p", params)

print("@@@@@@@@@@@@@@")

flops = FlopCountAnalysis(model, x)
print("FLOPs", flops.total())
print(parameter_count_table(model))

打印的结果为

[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.
[WARN] Cannot find rule for <class 'torch.nn.modules.container.Sequential'>. Treat it as zero Macs and zero Params.
[WARN] Cannot find rule for <class '__main__.net'>. Treat it as zero Macs and zero Params.
MACs 167616.0
p 576.0
@@@@@@@@@@@@@@
FLOPs 148992
| name            | #elements or shape   |
|:----------------|:---------------------|
| model           | 0.6K                 |
|  cnn            |  0.6K                |
|   cnn.0         |   0.6K               |
|    cnn.0.weight |    (64, 1, 8, 1)     |
|    cnn.0.bias   |    (64,)             |

对应的WARN不用理会，因为这几个类本来就没有计算

其实计算量的判定还有MACCs，MADDs等方法，具体可以参考

CNN的参数量、计算量（FLOPs、MACs）与运行速度_Dr鹏的博客-CSDN博客_模型复杂度

本文参考

网络模型计算量评估_WTHunt的博客-CSDN博客_网络计算量