前言

在深度学习的研究中,神经网络有时候不光要求网络的精度,在调整模型或者精简模型的时候,需要知道模型的运算量有多大。一般来说模型的参数量能够大致反映模型量级,但是运算量才是我们真正需要知道的。本系列文章旨在介绍神经网络中各种网络层的运算量计算方法,以及代码实现自动计算网络模型运算量。
神经网络的计算量主要体现在卷积层、全连接层,其它的激活层、BatchNorm层、池化层等占比重较少。

运算量计算方法

以下计算方法均针对batch size为1的情况,若batch size不为1,则乘上batch size即可。

1、卷积层

一般来说,采用了滑动窗口实现的卷积操作忽略了非线性计算开销,其FLOPs计算公式为:
神经网络批处理 神经网络flop_神经网络跳跃连接的作用
其中:神经网络批处理 神经网络flop_卷积_02分别代表卷积层输出的高和宽,神经网络批处理 神经网络flop_卷积_03为输入的通道数,神经网络批处理 神经网络flop_神经网络_04为卷积核size,神经网络批处理 神经网络flop_卷积_05为输出的通道数,神经网络批处理 神经网络flop_卷积_06为组卷积中的分组数groups,神经网络批处理 神经网络flop_运算量_07是针对有bias的卷积。
具体解释:

  • 对于输出tensor中的每一个数值,都对应着一次卷积操作,而这一次卷积操作是卷积核对于输入所有通道的一个局部区域进行的计算,也就是神经网络批处理 神经网络flop_运算量_08的运算量(加法比乘法少1次)
  • 输出一共神经网络批处理 神经网络flop_运算量_09个值,所以运算量为神经网络批处理 神经网络flop_神经网络_10
  • 对于组卷积,输入按通道分为神经网络批处理 神经网络flop_卷积_11组,针对每一组进行卷积操作,每组得到的结果组成最终输出,所以卷积操作运算量为:神经网络批处理 神经网络flop_运算量_12
  • 如果有bias的话,卷积操作过后,每个数值还需要加上bias,即神经网络批处理 神经网络flop_运算量_09次加法运算

2、全连接层

全连接层FLOPs计算公式为:
神经网络批处理 神经网络flop_神经网络跳跃连接的作用_14
其中神经网络批处理 神经网络flop_神经网络跳跃连接的作用_15分别为全连接层的输入和输出维数,神经网络批处理 神经网络flop_运算量_07针对bias。
这个应该很好理解,全连接层可以看作是向量和矩阵的乘法,同样是加法比乘法少1次。

3、池化层

池化层可以分为全局池化和一般池化两种,其FLOPs计算公式可以写为:
神经网络批处理 神经网络flop_神经网络_17
其中每个变量代表的含义与上述卷积层计算方法中一样。
具体解释:

  • 全局池化:针对输入所有值进行一次池化操作,不论是max、sum还是avg,都只需要对每个值算一次。
  • 一般池化:对于输出的每一个值,都是一次池化操作,每次池化操作的计算量就是池化过滤器参数的数量。

4、激活层

以使用最多的ReLU为例:
神经网络批处理 神经网络flop_GFLOPs_18
可见,激活层的运算量只跟输入大小有关,少数激活函数的运算量与输入为线性关系,基本上可以认为等于输入的大小。

5、BN层

对于BN来说,一般标配是conv + bn + relu,在上线使用过程中,一般会将bn和卷积层合并,相当于没了bn这一层,变成 conv +relu,所以bn其实不用考虑。

到此,已经介绍了神经网络中主要运算量的计算方式,后续将会根据上述计算方法实现对神经网络运算量的自动计算。
部分不常用的网络层没有介绍计算方法,之后有时间会加上,以上为个人理解,欢迎批评指正。