转载Unet网络参数的计算

转自：https://bbs.cvmart.net/articles/1632

1. 理论公式部分
   卷积层：

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

BN层：
，其中Ci为输入的channel数
（BN层有两个需要学习的参数，平移因子和缩放因子）

全连接层：
，Ti为输入向量的长度，To为输出向量的长度，其中第二项为偏置项参数量。 (不过目前全连接层已经逐渐被Global Average Pooling层取代了)

2. 实践部分
   我们首先摆出经典的UNet结构图。
   

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

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

这样定义的UNet主要有四个组件，3×3Conv，1×1Conv，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：

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

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