SAN (CVPR 2019) :基于二阶通道注意力机制的单图像超分网络

SAN:Second-Order Attention Network for Single Image Super-Resolution

论文地址：Second-Order Attention Network for Single Image Super-Resolution

代码地址：daitao/SAN: Second-order Attention Network for Single Image Super-resolution (CVPR-2019)

简介

​ 提出一种基于二阶统计信息的通道注意力机制，产生更好的表征能力，同时，模型对non-local机制也进行了优化，针对low-level任务，直接将non-local应用在整个图会导致计算量过大，于是采用了patch进行region-level的non-local机制。

现阶段问题

1. 现阶段的基于CNN的方法大多关注在如何设计更宽更深的网络，忽略了探索中间层的特征相关性，阻碍了CNN的表示能力。

2. SENet中的Channel Attention只关注了一阶统计量(eg. 全局池化)，忽略了高于一阶的统计量，阻碍了网络的判别能力

主要贡献

1. 设计了一种新型的可训练的二阶通道注意(SOCA)模块，通过使用二阶特征统计量进行更具鉴别性的表示，自适应地重新缩放通道特征
   • 二阶统计量比一阶统计量更具有鉴别性的表示（见1,2）
   • 协方差归一化(见1,3)对更具辨别力的表示起着至关重要的作用
2. 采用RL-NL module：采用region level的非局部non-local操作：不仅可以捕获长距离空间上下文信息，还可以扩大感受野
   • 对于low-level任务，链接中论文经过实验表明，适当的邻域的非局部操作表现的比全局non-local表现的更好。
   • 当特征图很大的时候，使用vanilla non-local计算负担也会很大。

方法概述

Second-order Channel Attention (SOCA)

​ 对于给定输入，将其特征reshape为$XwithC\times S，wheres=WH$,计算样本的协方差矩阵
$$\begin{array}{rl}\Sigma =& \mathrm{X}\bar{\mathrm{I}}{\mathrm{X}}^T,\mathrm{where}\bar{\mathbf{I}}=\frac{1}{s}(\mathbf{I}-\frac{1}{s}1)\\ & \mathbf{I}\mathrm{and}1\mathrm{are}\mathrm{the}s\times s\mathrm{identity}\mathrm{matrix}\mathrm{and}\mathrm{matrix}\mathrm{of}\mathrm{all}\mathrm{ones}\end{array}$$

​ 由于协方差归一化能提高模型的辨别性的表征能力，因此对$\Sigma$进行归一化。而因为$\Sigma$是对称半正定的矩阵，其具有特征值分解(EIG)如下：
$$\Sigma =U\Lambda U^T$$
$U$是一个正交矩阵，$\Lambda =diag(\lambda 1，\cdot \cdot \cdot ，\lambda C)$是具有非递增阶特征值的对角矩阵,协方差归一化可以转为：
$$\hat{Y}={\Sigma }^{\alpha }=U{\Lambda }^{\alpha }{U}^T$$
当$\alpha <1$，会非线性的缩小特征值大于1的值，并放大那些小于1的值。在贡献的[1]参考文献中表示，$\alpha =0.5$具有最好的表征能力。

计算协方差矩阵

class Covpool(Function):"""Global Covariance pooling layer"""@staticmethoddef forward(ctx, input):x = inputbatchSize = x.data.shape[0]# hwcdim = x.data.shape[1]h = x.data.shape[2]w = x.data.shape[3]# sM = h * w# Σ = X I_hat X^T,而I为SxS的矩阵,所以x需要reshape为dim,Mx = x.reshape(batchSize, dim, M)# I_hat=1/s(I-1/s 1)=(-1/s/s)*1+1/s*I,I and 1 are the s × s identity matrix and matrix of all onesI_hat = (1. / M) * torch.eye(M, M, device=x.device)+(-1. / M / M) * torch.ones(M, M, device=x.device)# 将I_hat转到和x的shape一样，因为存在batch，所以需要repeatI_hat = I_hat.view(1, M, M).repeat(batchSize, 1, 1).type(x.dtype)"""计算协方差矩阵Σ = X I_hat X^T"""# y = x I_hat x^T# x的shape为b,c,m,所以transpose是2，3维度# x.bmm(I_hat) 表示 x 和 I_hat 的批量矩阵乘法y = x.bmm(I_hat).bmm(x.transpose(1, 2))# 用于反向传播ctx.save_for_backward(input, I_hat)return y@staticmethoddef backward(ctx, grad_output):input, I_hat = ctx.saved_tensorsx = inputbatchSize = x.data.shape[0]dim = x.data.shape[1]h = x.data.shape[2]w = x.data.shape[3]M = h * wx = x.reshape(batchSize, dim, M)grad_input = grad_output + grad_output.transpose(1, 2)grad_input = grad_input.bmm(x).bmm(I_hat)grad_input = grad_input.reshape(batchSize, dim, h, w)return grad_input

基于Newton-Schulz迭代的快速矩阵归一化方法

Towards Faster Training of Global Covariance Pooling Networks by Iterative Matrix Square Root Normalization受到这篇论文的启发，文章中利用了Newton-Schulz迭代来加速协方差归一化的计算。对于${\Sigma }^{1/2}=U{\Lambda }^{1/2}{U}^T$，通过令$Y_0=\Sigma ,Z_0=I$,交替迭代更新如下：
$$\begin{array}{rl}{\mathbf{Y}}_n& =\frac{1}{2}{\mathbf{Y}}_{n-1}(3\mathbf{I}-{\mathbf{Z}}_{n-1}{\mathbf{Y}}_{n-1}),\\ {\mathbf{Z}}_n& =\frac{1}{2}(3\mathbf{I}-{\mathbf{Z}}_{n-1}{\mathbf{Y}}_{n-1}){\mathbf{Z}}_{n-1}.)\end{array}$$

由于Newton-Schulz迭代只局部收敛，为了保证收敛性 ，首先对$\Sigma$进行pre-norm归一化
$$\hat{\Sigma }=\frac{1}{tr(\Sigma )}\Sigma$$
其中$tr(\Sigma )={\sum }_i^C{\lambda }_i$表示$\Sigma$的迹。在这种情况下，能推断出$||\Sigma -I|{|}_2$等于$(\Sigma -I)$最大奇异值。$1-\frac{{\lambda }_i}{\sum i{\lambda }_i}$小于1，满足收敛条件.

再迭代之后，采用后补偿法，补偿在pre-norm中引起的数值波动，最后得到归一化协方差矩阵
$$\hat{Y}=\sqrt{tr(\Sigma )}{Y}_N,Nisfinaliter$$

class Sqrtm(Function):@staticmethoddef forward(ctx, input, iterN):x = inputbatchSize = x.data.shape[0]dim = x.data.shape[1]dtype = x.dtype# 3II3 = 3.0 * torch.eye(dim, dim, device=x.device).view(1, dim, dim).repeat(batchSize, 1, 1).type(dtype)# 计算tr(\Sigma)，乘以单位对角阵然后求和normA = (1.0 / 3.0) * x.mul(I3).sum(dim=1).sum(dim=1)# pre_normA = x.div(normA.view(batchSize, 1, 1).expand_as(x))# 让Y,Z具有相应的输出尺寸大小Y = torch.zeros(batchSize, iterN, dim, dim, requires_grad=False, device=x.device)Z = torch.eye(dim, dim, device=x.device).view(1, dim, dim).repeat(batchSize, iterN, 1, 1)if iterN < 2:ZY = 0.5 * (I3 - A)Y[:, 0, :, :] = A.bmm(ZY)else:"""iter1"""# 0.5(3I-Z_N-1Y_N-1)ZY = 0.5 * (I3 - A)# Y_1=0.5Y_0(3I-Z_0Y_0)=0.5A*(I3-A)Y[:, 0, :, :] = A.bmm(ZY)Z[:, 0, :, :] = ZYfor i in range(1, iterN - 1):# 3I-Z_N-1 Z_Y-1ZY = 0.5 * (I3 - Z[:, i - 1, :, :].bmm(Y[:, i - 1, :, :]))Y[:, i, :, :] = Y[:, i - 1, :, :].bmm(ZY)Z[:, i, :, :] = ZY.bmm(Z[:, i - 1, :, :])#最后一次迭代不用更新Z，直接求YZY = 0.5 * Y[:, iterN - 2, :, :].bmm(I3 - Z[:, iterN - 2, :, :].bmm(Y[:, iterN - 2, :, :]))# y_hat=\sqrt{ tr(\Sigma) } Y_N，后补偿y = ZY * torch.sqrt(normA).view(batchSize, 1, 1).expand_as(x)ctx.save_for_backward(input, A, ZY, normA, Y, Z)ctx.iterN = iterNreturn y@staticmethoddef backward(ctx, grad_output):input, A, ZY, normA, Y, Z = ctx.saved_tensorsiterN = ctx.iterNx = inputbatchSize = x.data.shape[0]dim = x.data.shape[1]dtype = x.dtypeder_postCom = grad_output * torch.sqrt(normA).view(batchSize, 1, 1).expand_as(x)der_postComAux = (grad_output * ZY).sum(dim=1).sum(dim=1).div(2 * torch.sqrt(normA))I3 = 3.0 * torch.eye(dim, dim, device=x.device).view(1, dim, dim).repeat(batchSize, 1, 1).type(dtype)if iterN < 2:der_NSiter = 0.5 * (der_postCom.bmm(I3 - A) - A.bmm(der_sacleTrace))else:dldY = 0.5 * (der_postCom.bmm(I3 - Y[:, iterN - 2, :, :].bmm(Z[:, iterN - 2, :, :])) -Z[:, iterN - 2, :, :].bmm(Y[:, iterN - 2, :, :]).bmm(der_postCom))dldZ = -0.5 * Y[:, iterN - 2, :, :].bmm(der_postCom).bmm(Y[:, iterN - 2, :, :])for i in range(iterN - 3, -1, -1):YZ = I3 - Y[:, i, :, :].bmm(Z[:, i, :, :])ZY = Z[:, i, :, :].bmm(Y[:, i, :, :])dldY_ = 0.5 * (dldY.bmm(YZ) -Z[:, i, :, :].bmm(dldZ).bmm(Z[:, i, :, :]) -ZY.bmm(dldY))dldZ_ = 0.5 * (YZ.bmm(dldZ) -Y[:, i, :, :].bmm(dldY).bmm(Y[:, i, :, :]) -dldZ.bmm(ZY))dldY = dldY_dldZ = dldZ_der_NSiter = 0.5 * (dldY.bmm(I3 - A) - dldZ - A.bmm(dldY))grad_input = der_NSiter.div(normA.view(batchSize, 1, 1).expand_as(x))grad_aux = der_NSiter.mul(x).sum(dim=1).sum(dim=1)for i in range(batchSize):grad_input[i, :, :] += (der_postComAux[i] \- grad_aux[i] / (normA[i] * normA[i])) \* torch.ones(dim, device=x.device).diag()return grad_input, None

region-level non local

​ 原始non-local机制参见Non-local Neural Networks，由于原始的是global 的non-local，当特征图较大，会导致计算量复杂；又根据经验表明，在合适的局部大小进行非局部操作能很好的适合low-level任务。因此论文采用了region-level non local。

​ 将图片切成四块，每一块中进行region-level non-local机制，最后在拼接在一起。

class Nonlocal_CA(nn.Module):def __init__(self, in_feat=64, inter_feat=32, reduction=8,sub_sample=False, bn_layer=True):super(Nonlocal_CA, self).__init__()# second-order channel attentionself.soca=SOCA(in_feat, reduction=reduction)# nonlocal moduleself.non_local = (NONLocalBlock2D(in_channels=in_feat,inter_channels=inter_feat, sub_sample=sub_sample,bn_layer=bn_layer))self.sigmoid = nn.Sigmoid()def forward(self,x):## divide feature map into 4 partbatch_size,C,H,W = x.shapeH1 = int(H / 2)W1 = int(W / 2)nonlocal_feat = torch.zeros_like(x)feat_sub_lu = x[:, :, :H1, :W1]feat_sub_ld = x[:, :, H1:, :W1]feat_sub_ru = x[:, :, :H1, W1:]feat_sub_rd = x[:, :, H1:, W1:]nonlocal_lu = self.non_local(feat_sub_lu)nonlocal_ld = self.non_local(feat_sub_ld)nonlocal_ru = self.non_local(feat_sub_ru)nonlocal_rd = self.non_local(feat_sub_rd)nonlocal_feat[:, :, :H1, :W1] = nonlocal_lunonlocal_feat[:, :, H1:, :W1] = nonlocal_ldnonlocal_feat[:, :, :H1, W1:] = nonlocal_runonlocal_feat[:, :, H1:, W1:] = nonlocal_rdreturn  nonlocal_feat

vanilla non-local机制如下：

class _NonLocalBlockND(nn.Module):def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded_gaussian',sub_sample=True, bn_layer=True):super(_NonLocalBlockND, self).__init__()assert dimension in [1, 2, 3]assert mode in ['embedded_gaussian', 'gaussian', 'dot_product', 'concatenation']# print('Dimension: %d, mode: %s' % (dimension, mode))self.mode = modeself.dimension = dimensionself.sub_sample = sub_sampleself.in_channels = in_channelsself.inter_channels = inter_channelsif self.inter_channels is None:self.inter_channels = in_channels // 2if self.inter_channels == 0:self.inter_channels = 1if dimension == 3:conv_nd = nn.Conv3dmax_pool = nn.MaxPool3dbn = nn.BatchNorm3delif dimension == 2:conv_nd = nn.Conv2dmax_pool = nn.MaxPool2dsub_sample = nn.Upsamplebn = nn.BatchNorm2delse:conv_nd = nn.Conv1dmax_pool = nn.MaxPool1dbn = nn.BatchNorm1dself.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,kernel_size=1, stride=1, padding=0)if bn_layer:self.W = nn.Sequential(conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,kernel_size=1, stride=1, padding=0),bn(self.in_channels))nn.init.constant_(self.W[1].weight, 0)nn.init.constant_(self.W[1].bias, 0)else:self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,kernel_size=1, stride=1, padding=0)nn.init.constant_(self.W.weight, 0)nn.init.constant_(self.W.bias, 0)self.theta = Noneself.phi = Noneself.concat_project = None# self.fc = nn.Linear(64,2304,bias=True)# self.sub_bilinear = nn.Upsample(size=(48,48),mode='bilinear')# self.sub_maxpool = nn.AdaptiveMaxPool2d(output_size=(48,48))if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,kernel_size=1, stride=1, padding=0)self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,kernel_size=1, stride=1, padding=0)if mode == 'embedded_gaussian':self.operation_function = self._embedded_gaussianelif mode == 'dot_product':self.operation_function = self._dot_productelif mode == 'concatenation':self.operation_function = self._concatenationself.concat_project = nn.Sequential(nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),nn.ReLU())elif mode == 'gaussian':self.operation_function = self._gaussianif sub_sample:self.g = nn.Sequential(self.g, max_pool(kernel_size=2))if self.phi is None:self.phi = max_pool(kernel_size=2)else:self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))def forward(self, x):''':param x: (b, c, t, h, w):return:'''output = self.operation_function(x)return outputdef _embedded_gaussian(self, x):batch_size,C,H,W = x.shape# x_sub = self.sub_bilinear(x) # bilinear downsample# x_sub = self.sub_maxpool(x) # maxpool downsample### g_x = x.view(batch_size, self.inter_channels, -1)# g_x = g_x.permute(0, 2, 1)## # theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)# # phi  =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)# # f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)# theta_x = x.view(batch_size, self.inter_channels, -1)# theta_x = theta_x.permute(0, 2, 1)# fc = self.fc(theta_x)# # phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)# # f = torch.matmul(theta_x, phi_x)# # return f# # f_div_C = F.softmax(fc, dim=-1)# return fc### g=>(b, c, t, h, w)->(b, 0.5c, t, h, w)->(b, thw, 0.5c)g_x = self.g(x).view(batch_size, self.inter_channels, -1)g_x = g_x.permute(0, 2, 1)# theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)# phi  =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)# f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)theta_x = theta_x.permute(0, 2, 1)phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)f = torch.matmul(theta_x, phi_x)# return ff_div_C = F.softmax(f, dim=-1)# return f_div_C# (b, thw, thw)dot(b, thw, 0.5c) = (b, thw, 0.5c)->(b, 0.5c, t, h, w)->(b, c, t, h, w)y = torch.matmul(f_div_C, g_x)y = y.permute(0, 2, 1).contiguous()y = y.view(batch_size, self.inter_channels, *x.size()[2:])W_y = self.W(y)z = W_y + xreturn zdef _gaussian(self, x):batch_size = x.size(0)g_x = self.g(x).view(batch_size, self.inter_channels, -1)g_x = g_x.permute(0, 2, 1)theta_x = x.view(batch_size, self.in_channels, -1)theta_x = theta_x.permute(0, 2, 1)if self.sub_sample:phi_x = self.phi(x).view(batch_size, self.in_channels, -1)else:phi_x = x.view(batch_size, self.in_channels, -1)f = torch.matmul(theta_x, phi_x)f_div_C = F.softmax(f, dim=-1)y = torch.matmul(f_div_C, g_x)y = y.permute(0, 2, 1).contiguous()y = y.view(batch_size, self.inter_channels, *x.size()[2:])W_y = self.W(y)z = W_y + xreturn zdef _dot_product(self, x):batch_size = x.size(0)g_x = self.g(x).view(batch_size, self.inter_channels, -1)g_x = g_x.permute(0, 2, 1)theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)theta_x = theta_x.permute(0, 2, 1)phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)f = torch.matmul(theta_x, phi_x)N = f.size(-1)f_div_C = f / Ny = torch.matmul(f_div_C, g_x)y = y.permute(0, 2, 1).contiguous()y = y.view(batch_size, self.inter_channels, *x.size()[2:])W_y = self.W(y)z = W_y + xreturn zdef _concatenation(self, x):batch_size = x.size(0)g_x = self.g(x).view(batch_size, self.inter_channels, -1)g_x = g_x.permute(0, 2, 1)# (b, c, N, 1)theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)# (b, c, 1, N)phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)h = theta_x.size(2)w = phi_x.size(3)theta_x = theta_x.repeat(1, 1, 1, w)phi_x = phi_x.repeat(1, 1, h, 1)concat_feature = torch.cat([theta_x, phi_x], dim=1)f = self.concat_project(concat_feature)b, _, h, w = f.size()f = f.view(b, h, w)N = f.size(-1)f_div_C = f / Ny = torch.matmul(f_div_C, g_x)y = y.permute(0, 2, 1).contiguous()y = y.view(batch_size, self.inter_channels, *x.size()[2:])W_y = self.W(y)z = W_y + xreturn zclass NONLocalBlock1D(_NonLocalBlockND):def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=True, bn_layer=True):super(NONLocalBlock1D, self).__init__(in_channels,inter_channels=inter_channels,dimension=1, mode=mode,sub_sample=sub_sample,bn_layer=bn_layer)class NONLocalBlock2D(_NonLocalBlockND):def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=True, bn_layer=True):super(NONLocalBlock2D, self).__init__(in_channels,inter_channels=inter_channels,dimension=2, mode=mode,sub_sample=sub_sample,bn_layer=bn_layer)

结论

我们提出了一个深度二阶注意力网络 (SAN) 来实现准确的图像 SR。具体来说，非局部增强残差组 (NLRG) 结构允许 SAN 通过在网络中嵌入非局部操作来捕获长距离依赖和结构信息。同时，NLRG 允许通过共享源跳跃连接绕过 LR 图像中丰富的低频信息。除了利用空间特征相关性外，我们还提出了二阶通道注意(SOCA)模块，通过全局协方差池化来学习特征相互依赖性，以获得更具鉴别性的表示。在 BI 和 BD 退化模型的 SR 上的大量实验表明，我们的 SAN 在定量和视觉结果方面的有效性。