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

1、公式

  以下公式适用于各种情况的卷积层,如普通卷积、膨胀卷积、分组卷积、可分离卷积等:

参数量:
  无bias时: k H × k W × C i n / g × C o u t k_{H }×k_{W }×C_{in}/g×C_{out} kH×kW×Cin/g×Cout
  有bias时: ( k H × k W × C i n / g + 1 ) × C o u t (k_{H }×k_{W }×C_{in}/g + 1)×C_{out} (kH×kW×Cin/g+1)×Cout

浮点计算数FLOPs(指所有的乘法和加法运算):
  无bias时: ( 2 × k H × k W × C i n / g − 1 ) × C o u t × H o u t × W o u t (2×k_{H }×k_{W }×C_{in}/g - 1)×C_{out}×H_{out}×W_{out} (2×kH×kW×Cin/g1)×Cout×Hout×Wout
  有bias时: 2 × k H × k W × C i n / g × C o u t × H o u t × W o u t 2×k_{H }×k_{W }×C_{in}/g×C_{out}×H_{out}×W_{out} 2×kH×kW×Cin/g×Cout×Hout×Wout

式中各量对应的pytorch卷积层的配置参数如下:
torch.nn.Conv2d(in_channels= C i n C_{in } Cin, out_channels= C o u t C_{out } Cout, kernel_size= ( k H , k W ) (k_{H },k_{W }) (kH,kW), groups=g)

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

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

2、代码

  如果只计算网络的参数量,不必用上述公式,直接从model中提取即可:

params = sum([v.numel() for k,v in model.state_dict().items()])
#计算其中部分网络的参数量
params_bb = sum([v.numel() for k,v in model.state_dict().items() if k.startswith('backbone.')])

如果计算所有卷积层的参数量,则可以

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操作得到各中间层的特征图尺寸并代入上面的公式计算。代码如下,注意此代码只计算网络中的卷积层,不计算全连接层等其他层,代码也同时计算了卷积层的参数量:

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)
Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐