目录
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的读内存是输入占内存大小+ 中间网络参数的大小
写内存是输出占内存大小