聊聊神经网络结构以及计算量和内存使用

目录

1. 前言

2. torchsummary

3. torchstat

3.1 Conv 层计算FLOPs和MAC

3.2 ReLU 计算FLOPs和MAC

3.3 MaxPool 计算FLOPs和MAC

3.4 fc 计算FLOPs和MAC

4. summary

---

1. 前言

最近在看轻量级网络的paper，因此来简单聊聊神经网络计算量和使用内存的情况

这里只计算两个参数FLOPs 和 MAC

FLOPs是神经网络执行一次前向传播的计算量，也就是乘法-加法浮点运算次数

MAC ：memory access cost ，网络参数的内存消耗

FLOPS，最后一个大写S代表的是硬件每秒的浮点运算次数

一般来说，对卷积的计算方式如下：

具体的计算再后面讲解

2. torchsummary

torchsummary 模块可以查看网络的架构

下载的方式就是：pip install torchsummary

如图，搭建了一个simple的net

这里使用torchsummary查看，这里需要注意的是，summary的input默认是在cuda上的，所以这里要将网络放到cuda上，或者可以利用device参数将网络和input都放到cpu上

from torchsummary import summary


model = TestNet()
print(summary(model.cuda(),input_size=(3,24,24)))

打印结果：

这里关于如何计算网络参数不在介绍，可以参考----> 如何计算神经网络参数的个数？

下面只会做简单介绍

3. torchstat

安装也是：pip install torchstat

这里测试代码如下：

from torchstat import stat

model = TestNet()
print(stat(model,input_size=(3,24,24)))

打印如下：

 这里的layer0，1，2 就是定义网络的对应操作

3.1 Conv 层计算FLOPs和MAC

如代码所示，这里的输入是3*24*24的彩色图像，经过3*3 的same 卷积得到6*24*24的特征图

所以这里的参数量是 3*3*3*6 + 6(bias) = 168

FLOPs和MAC参考下面的公式（下图没有加上bias）

FLOPs = 输入图像的size * 卷积核的size * 输入个数(input channel) * 输出个数(output channel) + 输入图像的size * 卷积核的个数(输出的channel个数)

这里要注意，第二个加法是bias的计算量，因为图像的每一个像素点卷积后都要加一个偏置

所以这里的FLOPs = (24*24)*(3*3)*(3)*(6)+(24*24)*(6) =  96768 ,这里的括号是为了和红色的公式对应

而MAC再torchstat是分开来表示的，也就是读内存MemRead(B) 和 写内存MemWrite(B)，两者相加就是计算的MAC，即MemR+W(B)

读内存 =  输入图像的size * 输入图像的个数(input channel) + 卷积核的size * 输入个数(input channel)* 输出个数(out channel) + 输出图像的个数(out channel) ---> bias

bias虽然计算量复杂，但其实就是一个标量，然后有多少个卷积核，就有多少个bias，也就有多少个内存占据

写内存 =  输出图像的size * 输出图像的个数(out channel) 

so ，

读内存MemRead(B) = (24*24)*(3) + (3*3)*(3)*(6) + (6) = 1896,这里是float32，所以乘上4就等于7584

写内存MemWrite(B) = 24*24*6 = 3456，乘上4就是 13824

因此Conv 卷积的总内存cost = 7584+13824 = 21408

3.2 ReLU 计算FLOPs和MAC

ReLU层的信息在第二个

 因为ReLU 是没有参数的，是elementwise，所以ReLU的计算量就是上一层输出的参数

也就是 6 *24 *24 = 3456，而读内存就是上一层的参数个数3456 * 4 = 13824，写内存也就是上一层的参数经过ReLU的结果，也就是3456 * 4 = 13824，结果不变

3.3 MaxPool 计算FLOPs和MAC

max pooling 层的信息在第三个

计算量和FLOPs和ReLU一样，有所区别的是，max的输出是输入的1/4，也就是这里的写内存是输入的1/4

3.4 fc 计算FLOPs和MAC

如下：

其实这里和conv 层一样了

参数量 = input * output + output = 864*120+120 = 103800

计算量FLOPs = input * output + output = 864*120+120 = 103800

读取内存 = input * output + output + input = 864*120+120 + 864 = 104664，×4 = 418656

写入内存 = output = 120，×4 = 480

4. summary

关于参数：

计算参数很简单，卷积层就是计算卷积核的个数，然后每个卷积核包含的参数，有偏置的话，就加上卷积核的个数。全连接层就是输入 * 输出，有偏置，就加上输出的个数

关于计算量FLOPs：

卷积层的计算量不仅仅包括卷积核本身的计算，而且要根据输入图像的size进行计算。bias类似于一个逐元素计算，所以也和输入图像的size有关

关于MAC：

就是看每个层的元素个数即可，float的话，每个元素占据4个字节

总共分为三个部分：输入占内存大小 + 输出占内存大小 + 中间网络参数的大小

而stat的读内存是输入占内存大小+ 中间网络参数的大小

写内存是输出占内存大小