CondConv代码解析

论文地址：https://arxiv.org/abs/1904.04971
代码（Pytorch版本，网友复现）：https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/model/conv/CondConv.py

最近几个公众号对 CondConv 这个经典的方法进行了回顾，我们也对该方法的代码进行解析。

CondConv 是谷歌团队发表于NeurIPS 2019的一个工作，非常经典，可能看作为动态网络的开端，最重要的是原理简单，效果非常好。

卷积层设计中的一个基本假设是，相同的卷积核应用于数据集中的每个样本 。在本文中，作者提出的条件参数化卷积（CondConv） ，通过基于输入来动态计算卷积核，从而避免了传统静态的卷积中所有样本共享一个卷积核的缺点。具体地，作者将CondConv层中的卷积核参数化为n个专家（卷积核）的线性组合，即\((\alpha_1 W_1+\alpha_2 W_2+\ldots+\alpha_n W_n)\)，其中\(W_i\)是卷积核，\(\alpha_i\)是通过反向传播学习的参数。

该方法的总体框架如上图所示，下面结合网友代码分析细节，为了方便理解 ，把原代码中的 initial_weights 和 bias 相关的部分删掉了。首先，是 rount_fn 这部分，是 attention 函数，输入为 [N, C, H, W]，需要两个参数，in_planes为输特征通道数，K 为专家个数。输出为[N, K]，即卷积核的权重。

# 输入为 [N, C, H, W]，需要两个参数，in_planes为输特征通道数，K 为专家个数
class Attention(nn.Module):
    def __init__(self,in_planes,K):
        super().__init__()
        self.avgpool=nn.AdaptiveAvgPool2d(1)
        self.net=nn.Conv2d(in_planes, K, kernel_size=1)
        self.sigmoid=nn.Sigmoid()

    def forward(self,x):
        # 将输入特征全局池化为 [N, C, 1, 1]
        att=self.avgpool(x)
        # 使用1X1卷积，转化为 [N, K, 1, 1]
        att=self.net(att)
        # 将特征转化为二维 [N, K]
        att=att.view(x.shape[0],-1) 
        # 使用 sigmoid 函数输出归一化到 [0,1] 区间
        return self.sigmoid(att)

下面是 CondConv 主函数，参数和普通卷积非常类似，只是多了 K，即专家个数。

class CondConv(nn.Module):
    def __init__(self,in_planes,out_planes,kernel_size,stride,padding=0,
                 groups=1,K=4):
        super().__init__()
        self.in_planes = in_planes
        self.out_planes = out_planes
        self.K = K
        self.groups = groups
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.attention = Attention(in_planes=in_planes,K=K)
        self.weight = nn.Parameter(torch.randn(K,out_planes,in_planes//groups,
                                             kernel_size,kernel_size),requires_grad=True)

    def forward(self,x):
        # 调用 attention 函数得到归一化的权重 [N, K]
        N,in_planels, H, W = x.shape
        softmax_att=self.attention(x)
        # 把输入特征由 [N, C_in, H, W] 转化为 [1, N*C_in, H, W]
        x=x.view(1, -1, H, W)

        # 生成随机 weight [K, C_out, C_in/groups, 3, 3] (卷积核一般为3*3)
        # 注意添加了 requires_grad=True，这样里面的参数是可以优化的
        weight = self.weight
        # 改变 weight 形状为 [K, C_out*(C_in/groups)*3*3]
        weight = weight.view(self.K, -1) 

        # 矩阵相乘：[N, K] X [K, C_out*(C_in/groups)*3*3] = [N, C_out*(C_in/groups)*3*3]
        aggregate_weight = torch.mm(softmax_att,weight)
        # 改变形状为：[N*C_out, C_in/groups, 3, 3]，即新的卷积核权重
        aggregate_weight = aggregate_weight.view(
            N*self.out_planes, self.in_planes//self.groups,
            self.kernel_size, self.kernel_size)
        # 用新生成的卷积核进行卷积，输出为 [1, N*C_out, H, W]
        output=F.conv2d(x,weight=aggregate_weight,
                        stride=self.stride, padding=self.padding,
                        groups=self.groups*N)
        # 形状恢复为 [N, C_out, H, W]        
        output=output.view(N, self.out_planes, H, W)
        return output

CondConv 开始引进入动态卷积，增加卷积核生成函数的大小和复杂性 。本质上是把注意力机制应用在卷积核上了。CondConv还强调了一个重要问题，即如利用样本的特点来提高模型性能。