查看原文
其他

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

Hello王叔叔 极市平台 2021-09-19

加入极市专业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、图像分割等极市技术交流群,更有每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、行业技术交流一起来让思想之光照的更远吧~


△长按添加极市小助手


△长按关注极市平台,获取最新CV干货


觉得有用麻烦给个在看啦~  

: . Video Mini Program Like ,轻点两下取消赞 Wow ,轻点两下取消在看

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存