神经网络参数量的计算：以UNet为例

2020 年 3 月 17 日 极市平台

加入极市专业CV交流群，与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流！

同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总，行业技术交流。关注 极市平台 公众号，回复 加群，立刻申请入群~

作者：Hello王叔叔

知乎链接：https://zhuanlan.zhihu.com/p/57437131

本文已由作者授权转载，未经允许，不得二次转载。

之前写回答【UNet为什么在医疗影像领域表现好】时提到了UNet的参数量，引发了一些疑问和讨论。

比如评论中提问者说UNet 是31M，还有一些算出28M的，而【UNet++论文】中提到的UNet是7.76M？所以就想写一篇文章来讲一下如何计算网络的参数量。为了搞清楚参数量的计算，首先我们先搞清楚参数量的定义。

参数量就是指，模型所有带参数的层的权重参数总量。视觉类网络组件中带参数的层，主要有：卷积层、BN层、全连接层等。（注意：激活函数层(relu等)和Maxpooling层、Upsample层是没有参数的，不需要学习，他们只是提供了一种非线性的变换）

1. 理论公式部分

卷积层

其中为卷积核大小，为输入channel数，为输出的channel数(也是filter的数量)，算式第二项是偏置项的参数量。(虽然一般不写偏置项，因为不会影响总参数量的数量级，但是我们为了准确起见，把偏置项的参数量也考虑进来）

BN层

，其中为输入的channel数。

BN层有两个需要学习的参数，平移因子和缩放因子。

全连接层

，为输入向量的长度，为输出向量的长度，其中第二项为偏置项参数量。

不过目前全连接层已经逐渐被Global Average Pooling层取代了。

有了公式，是时候干一番大事业了哈哈。我们将以UNet为例，来实践一下参数量的计算。

2. 实践部分

我们首先摆出经典的UNet结构图。

我们把UNet共分为5个Stage，分别计算每个stage的参数量。每个stage的filter数量为

[32,64,128,256,512]，相比于UNet原文，我们把UNet的channel数缩小了两倍，大多数论文也的确是这么做的。同时，我们设置UNet上采样方式为TransposeConv（转置卷积），并在每个 Conv后加入BN层。最后假定，原始输入channel为1，输出分割图为两类(含背景)，这样最终得到我们要计算参数量的UNet。

这样定义的UNet主要有四个组件， Conv， Conv，TransposeConv和BN层。

我们先计算Conv，再计算BN层。

Stage1：

Stage2：

Stage3:

Stage4:

Stage5:

TransposeConv：

目前为止，我们把以上所有Conv得到的参数量求和，得到了没有加BN的UNet参数量。

接下来，我们计算BN层的参数量，易得：

和刚求得的参数量求和，

加入了BN的UNet参数量为7765442。

我们得到的两种UNet的参数量，如何验证其正确性呢？或者说我们不可能每次都手动计算这些网络的参数量，这就需要脚本去帮我们去计算网络的参数量。代码如下：

def count_param(model):
    param_count = 0
    for param in model.parameters():
        param_count += param.view(-1).size()[0]
    return param_count

我们通过项目【UNet-family】中提供的脚本，可以直接使用命令python UNet.py来计算网络的参数量。通过设置参数is_batchnorm我们得到两次计算结果如下：

有BN：

没有BN：

赶快看一眼我们之前算的参数量，一模一样，说明我们算的没毛病。

有没有一种再也不怕面试官问自己参数量计算的感觉。

3. 拓展延伸

我把几种UNet变形的参数量列举如下，大家也可以自己验证。

（其中Feature Scale是指channel的scale相对于原始UNet降了几倍，如果Feature scale为2，则filter的数量为[32,64,128,256,512])

4. 留个作业可好

最后，留个小实践，大家可以去计算下UNet++的参数量，看和论文中是不是一致？

附上我自己计算的结果。

参考链接：

1. 卷积神经网络的复杂度分析，https://zhuanlan.zhihu.com/p/31575074

2. Batch Normalization原理与实战，https://zhuanlan.zhihu.com/p/34879333

-END -

推荐阅读：

极市平台视觉算法季度赛，提供真实应用场景数据和免费算力，特殊时期，一起在家打比赛吧！

添加极市小助手微信 （ID : cv-mart） ，备注： 研究方向-姓名-学校/公司-城市 （如：目标检测-小极-北大-深圳），即可申请加入 目标检测、目标跟踪、人脸、工业检测、医学影像、三维&SLAM、图像分割等极市技术交流群 ，更有 每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、 干货资讯汇总、行业技术交流 ，一起来让思想之光照的更远吧~