卷积参数量和计算量(FLOPs)的计算公式及代码

我们经常用参数量和浮点计算数FLOPs来衡量卷积网络的复杂度。下面推导其公式并在pytorch中实现，以二维卷积Conv2d为例。

1、公式

以下公式适用于各种情况的卷积层，如普通卷积、膨胀卷积、分组卷积、可分离卷积等：
参数量:
无bias时： $k_H\times k_W\times C_{in}/g\times C_{out}$
有bias时： $(k_H\times k_W\times C_{in}/g+1)\times C_{out}$
浮点计算数FLOPs（指所有的乘法和加法运算）：
无bias时： $(2\times k_H\times k_W\times C_{in}/g-1)\times C_{out}\times H_{out}\times W_{out}$
有bias时： $2\times k_H\times k_W\times C_{in}/g\times C_{out}\times H_{out}\times W_{out}$
式中各量对应的pytorch卷积层的配置参数如下：
torch.nn.Conv2d(in_channels=$C_{in}$, out_channels=$C_{out}$, kernel_size=$(k_H,k_W)$, groups=g)

对于可分离卷积，网上有的文章也单拎出来对其公式进行了推导，其实可分离卷积就是一个带groups参数的分组卷积加上一个1*1全卷积，不需要额外再推导，两部分都用上式计算之后再加起来就可以了。

推导过程也很简单。根据卷积的原理，
（1）卷积核的大小是$k_H\times k_W\times C_{in}$，卷积核的数量就是输出通道数$C_{out}$，无bias时，所有卷积核的元素数量就是参数量，乘起来就可以了。
（2）有bias时每个卷积核对应一个bias，所以在每个卷积核的参数量上再加1
（3）卷积计算时，无bias时，所有卷积核元素和对应位置的输入特征图张量一一相乘再相加，乘法次数就是卷积核元素数，注意相加的时候是逐个往第一个数上加的所以总加法次数是卷积核元素数减1，就是$(2\times k_H\times k_W\times C_{in}/g-1)$；而卷积核在一个位置的运算只能产生一个输出数据点，随着卷积核的平移（膨胀卷积的跳点平移也是一样）共产生了$C_{out}\times H_{out}\times W_{out}$个输出数据点，把这两部分乘起来就是总计算量。
（4）有bias时，每个输出数据点还要再加上bias，多一次加法，正好把上式的减1抵消掉了。
（5）分组卷积时，根据分组卷积原理，相当于进行了g个输入通道为$C_{in}/g$，输出通道为$C_{out}/g$的卷积，所以把上面所有式的$C_{in}$，$C_{out}$替换为$C_{in}/g$，$C_{out}/g$，再乘以g后抵消掉$C_{out}/g$的分母，就得到了最终公式。

2、代码

如果只计算所有卷积层的参数量，可以不用上述公式，直接从model中提取即可

total_para_nums = 0
for n,m in model.named_modules():
    if isinstance(m,nn.Conv2d):
        total_para_nums += m.weight.data.numel()
        if m.bias is not None:
            total_para_nums += m.bias.data.numel()
print('total parameters:',total_para_nums）

但FLOPs计算涉及到中间层特征图的尺寸，不能从model中得到，必须使用钩子hook操作。代码如下，此代码也同时计算参数量，可显示各层的参数量、FLOPs以及所有层的总和：

model = XXX_Net()#注：以下代码放在模型实例化之后，模型名用model

def my_hook(Module, input, output):
    outshapes.append(output.shape)
    modules.append(Module)

names,modules,outshapes = [],[],[]
for name,m in model.named_modules():
    if isinstance(m,nn.Conv2d):
        m.register_forward_hook(my_hook)
        names.append(name)

def calc_paras_flops(modules,outshapes):
    total_para_nums = 0
    total_flops = 0
    for i,m in enumerate(modules):
        Cin = m.in_channels
        Cout = m.out_channels
        k = m.kernel_size
        #p = m.padding
        #s = m.stride
        #d = m.dilation
        g = m.groups
        Hout = outshapes[i][2]
        Wout = outshapes[i][3]
        if m.bias is None:
            para_nums = k[0] * k[1] * Cin / g * Cout
            flops = (2 * k[0] * k[1] * Cin/g - 1) * Cout * Hout * Wout
        else:
            para_nums = (k[0] * k[1] * Cin / g +1) * Cout
            flops = 2 * k[0] * k[1] * Cin/g * Cout * Hout * Wout
        para_nums = int(para_nums)
        flops = int(flops)
        print(names[i], 'para:', para_nums, 'flops:',flops)
        total_para_nums += para_nums
        total_flops += flops
    print('total conv parameters:',total_para_nums, 'total conv FLOPs:',total_flops)
    return total_para_nums, total_flops

input = torch.rand(32,3,224,224)#需要先提供一个输入张量
y = model(input)
total_para_nums, total_flops = calc_paras_flops(modules,outshapes)