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还强调了一个重要问题,即如利用样本的特点来提高模型性能。