YOLOv5改进系列(8)——添加SOCA注意力机制
🚀一、SOCA介绍
- 论文题目:《Second-order Attention Network for Single Image Super-Resolution》
- 论文地址:CVPR19 超分辨率
- 代码地址:GitHub - daitao/SAN: Second-order Attention Network for Single Image Super-resolution (CVPR-2019)
1.1 简介
近年来,深度卷积神经网络(CNN)在单图像超分辨率(SISR)中得到了广泛的研究,并取得了显著的性能。然而,大多数现有的基于CNN的SISR方法主要侧重于更广泛或更深入的架构设计,而忽略了中间层的特征相关性,因此阻碍了CNN的代表能力。为了解决这个问题,论文作者提出了一种二阶注意力网络(SAN),用于更强大的特征表达和特征相关性学习。具体而言,开发了一种新的可训练二阶注意(SOCA)模块,通过使用二阶特征统计来自适应地重新调整信道方向特征,以获得更具鉴别性的表示。
1.2 SAN网络

从上图中可以看出SAN的主要由四部分组成:
- 浅层特征提取(shallow feature extraction)即第一个卷积
- 非局部增强残差组(NLRG) 提取深度特征(deep feature,DF)
- 上采样模块(upscale module)
- 重建模块(reconstruction part)即最后一个卷积
1.3.二阶通道注意力(SOCA)
以前大多数基于CNN的SR模型都没有考虑功能的相互依赖性。为了利用这些信息,在CNN中引入了SENet,以重新缩放图像SR的信道特征。然而,SENet仅通过全局平均池利用特征的一阶统计,而忽略高于一阶的统计,从而阻碍了网络的辨别能力。另一方面,最近的研究表明特征的二阶统计分布更有利于获得有区分度的表达,如此才诞生了SOCA。
二阶注意力机制(SOCA)能够更好地学习特征之间的联系,此模块通过利用二阶特征的分布自适应的学习特征的内部依赖关系,SOCA的机制是网络能够专注于更有益的信息且能够提高判别学习的能力。此外,原文提出了一种非局部加强残差组结构能进一步结合非局部操作来提取长程的空间上下文信息。通过堆叠非局部残差组,本文的方法能够利用LR图像的信息且能够忽略低频信息。
🚀二、在backbone末端添加SOCA注意力机制方法
2.1 添加顺序
(1)models/common.py --> 加入新增的网络结构
(2) models/yolo.py --> 设定网络结构的传参细节,将SOCA类名加入其中。(当新的自定义模块中存在输入输出维度时,要使用qw调整输出维度)
(3) models/yolov5*.yaml --> 新建一个文件夹,如yolov5s_SOCA.yaml,修改现有模型结构配置文件。(当引入新的层时,要修改后续的结构中的from参数)
(4) train.py --> 修改‘--cfg’默认参数,训练时指定模型结构配置文件
2.2 具体添加步骤
第①步:在common.py中添加SOCA模块
将下面的SOCA代码复制粘贴到common.py文件的末尾
- # SOCA moudle 单幅图像超分辨率
- from torch.autograd import Function
- class Covpool(Function):
- @staticmethod
- def forward(ctx, input):
- x = input
- batchSize = x.data.shape[0]
- dim = x.data.shape[1]
- h = x.data.shape[2]
- w = x.data.shape[3]
- M = h*w
- x = x.reshape(batchSize,dim,M)
- I_hat = (-1./M/M)*torch.ones(M,M,device = x.device) + (1./M)*torch.eye(M,M,device = x.device)
- I_hat = I_hat.view(1,M,M).repeat(batchSize,1,1).type(x.dtype)
- y = x.bmm(I_hat).bmm(x.transpose(1,2))
- ctx.save_for_backward(input,I_hat)
- return y
- @staticmethod
- def backward(ctx, grad_output):
- input,I_hat = ctx.saved_tensors
- x = input
- batchSize = x.data.shape[0]
- dim = x.data.shape[1]
- h = x.data.shape[2]
- w = x.data.shape[3]
- M = h*w
- x = 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
- class Sqrtm(Function):
- @staticmethod
- def forward(ctx, input, iterN):
- x = input
- batchSize = x.data.shape[0]
- dim = x.data.shape[1]
- dtype = x.dtype
- I3 = 3.0*torch.eye(dim,dim,device = x.device).view(1, dim, dim).repeat(batchSize,1,1).type(dtype)
- normA = (1.0/3.0)*x.mul(I3).sum(dim=1).sum(dim=1)
- A = x.div(normA.view(batchSize,1,1).expand_as(x))
- 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:
- ZY = 0.5*(I3 - A)
- Y[:,0,:,:] = A.bmm(ZY)
- Z[:,0,:,:] = ZY
- for i in range(1, iterN-1):
- ZY = 0.5*(I3 - Z[:,i-1,:,:].bmm(Y[:,i-1,:,:]))
- Y[:,i,:,:] = Y[:,i-1,:,:].bmm(ZY)
- Z[:,i,:,:] = ZY.bmm(Z[:,i-1,:,:])
- ZY = 0.5*Y[:,iterN-2,:,:].bmm(I3 - Z[:,iterN-2,:,:].bmm(Y[:,iterN-2,:,:]))
- y = ZY*torch.sqrt(normA).view(batchSize, 1, 1).expand_as(x)
- ctx.save_for_backward(input, A, ZY, normA, Y, Z)
- ctx.iterN = iterN
- return y
- @staticmethod
- def backward(ctx, grad_output, der_sacleTrace=None):
- input, A, ZY, normA, Y, Z = ctx.saved_tensors
- iterN = ctx.iterN
- x = input
- batchSize = x.data.shape[0]
- dim = x.data.shape[1]
- dtype = x.dtype
- der_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
- def CovpoolLayer(var):
- return Covpool.apply(var)
- def SqrtmLayer(var, iterN):
- return Sqrtm.apply(var, iterN)
- class SOCA(nn.Module):
- # second-order Channel attention
- def __init__(self, channel, reduction=8):
- super(SOCA, self).__init__()
- self.max_pool = nn.MaxPool2d(kernel_size=2)
- self.conv_du = nn.Sequential(
- nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
- nn.ReLU(inplace=True),
- nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
- nn.Sigmoid()
- )
- def forward(self, x):
- batch_size, C, h, w = x.shape # x: NxCxHxW
- N = int(h * w)
- min_h = min(h, w)
- h1 = 1000
- w1 = 1000
- if h < h1 and w < w1:
- x_sub = x
- elif h < h1 and w > w1:
- W = (w - w1) // 2
- x_sub = x[:, :, :, W:(W + w1)]
- elif w < w1 and h > h1:
- H = (h - h1) // 2
- x_sub = x[:, :, H:H + h1, :]
- else:
- H = (h - h1) // 2
- W = (w - w1) // 2
- x_sub = x[:, :, H:(H + h1), W:(W + w1)]
- cov_mat = CovpoolLayer(x_sub) # Global Covariance pooling layer
- cov_mat_sqrt = SqrtmLayer(cov_mat,5) # Matrix square root layer( including pre-norm,Newton-Schulz iter. and post-com. with 5 iteration)
- cov_mat_sum = torch.mean(cov_mat_sqrt,1)
- cov_mat_sum = cov_mat_sum.view(batch_size,C,1,1)
- y_cov = self.conv_du(cov_mat_sum)
- return y_cov*x
如下图所示:

第②步:在yolo.py文件里的parse_model函数加入类名
首先找到yolo.py里面parse_model函数的这一行

然后把刚才加入的类SOCA添加到这个注册表里面:

或者可以在下面的位置这样加,原理和上面是一样的:
- elif m is SOCA:
- c1, c2 = ch[f], args[0]
- if c2 != no:
- c2 = make_divisible(c2 * gw, 8)
- args = [c1, *args[1:]]

解释一下这段代码:
这段是一个判断语句,如果模块 m 在SOCA中,那么就将模块m对应的输入通道数和输出通道数的值分别赋值给 c1 和 c2,然后对 c2进行与之前相同的处理,接下来,将 c1、 c2 以及 args[1:] 作为元素,组成新的列表,作为更新后的 args。
第③步:创建自定义的yaml文件
首先在models文件夹下复制yolov5s.yaml 文件,粘贴并重命名为 yolov5s_SOCA.yaml

接着修改 yolov5s_SOCA.yaml ,将SOCA模块加到我们想添加的位置。
这里我先介绍第一种,第一种是将SOCA模块放在backbone部分的最末端,这样可以使注意力机制看到整个backbone部分的特征图,将具有全局视野,类似于一个小transformer结构。
将 [-1,1,SOCA,[1024]]添加到 SPPF 的下一层。即下图中所示位置:

同样的下面的head也得修改:

这里我们要把后面两个Concat的from系数分别由[ − 1 , 14 ] , [ − 1 , 10 ]改为[ − 1 , 15 ],[ − 1 , 11 ]。
然后将Detect原始的from系数[ 17 , 20 , 23 ]要改为[ 18 , 21 , 24 ] 。

yolov5s_SOCA.yaml 完整代码:
- # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
- # Parameters
- nc: 20 # number of classes
- depth_multiple: 0.33 # model depth multiple
- width_multiple: 0.50 # layer channel multiple
- anchors:
- - [10,13, 16,30, 33,23] # P3/8
- - [30,61, 62,45, 59,119] # P4/16
- - [116,90, 156,198, 373,326] # P5/32
- # YOLOv5 v6.0 backbone+SE
- backbone:
- # [from, number, module, args]
- [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
- [-1, 3, C3, [128]],
- [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
- [-1, 6, C3, [256]],
- [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
- [-1, 9, C3, [512]],
- [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
- [-1, 3, C3, [1024]],
- [-1, 1, SPPF, [1024, 5]], # 10
- [-1, 1, SOCA,[1024]],
- ]
- # YOLOv5 v6.1 head
- head:
- [[-1, 1, Conv, [512, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 6], 1, Concat, [1]], # cat backbone P4
- [-1, 3, C3, [512, False]], # 14
- [-1, 1, Conv, [256, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 4], 1, Concat, [1]], # cat backbone P3
- [-1, 3, C3, [256, False]], # 18 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]],
- [[-1, 15], 1, Concat, [1]], # cat head P4
- [-1, 3, C3, [512, False]], # 21 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]],
- [[-1, 11], 1, Concat, [1]], # cat head P5
- [-1, 3, C3, [1024, False]], # 24 (P5/32-large)
- [[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
第④步:验证是否加入成功
在yolo.py 文件里面配置改为我们刚才自定义的yolov5s_SOCA.yaml


然后我们运行yolo.py

这样就添加成功啦~
第⑤步:修改train.py中 ‘--cfg’默认参数
我们先找到 train.py 文件的parse_opt函数,然后将第二行‘--cfg’的 default改为yolov5s_SOCA.yaml,然后就可以开始训练啦~

🚀三、在C3后添加SOCA注意力机制方法
第二种是将SOCA放在backbone部分每个C3模块的后面,这样可以使注意力机制看到局部的特征,每层进行一次注意力,可以分担学习压力。
步骤和方法1相同,只是yaml文件不同。
所以接下来只放修改yaml文件的部分~
第③步:创建自定义的yaml文件
将SOCA模块放在每个C3模块的后面,要注意通道的变化。
如下图所示:

同样的,下面的head部分也要做相应的修改:

第二种方法的 yolov5s_SOCA.yaml 完整代码:
- # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
- # Parameters
- nc: 80 # number of classes
- depth_multiple: 0.33 # model depth multiple
- width_multiple: 0.50 # layer channel multiple
- anchors:
- - [10,13, 16,30, 33,23] # P3/8
- - [30,61, 62,45, 59,119] # P4/16
- - [116,90, 156,198, 373,326] # P5/32
- # YOLOv5 v6.0 backbone
- backbone:
- # [from, number, module, args]
- [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
- [-1, 3, C3, [128]],
- [-1, 3, SOCA, [128]],
- [-1, 1, Conv, [256, 3, 2]], # 4-P3/8
- [-1, 6, C3, [256]],
- [-1, 3, SOCA, [256]],
- [-1, 1, Conv, [512, 3, 2]], # 7-P4/16
- [-1, 9, C3, [512]],
- [-1, 3, SOCA, [512]],
- [-1, 1, Conv, [1024, 3, 2]], # 10-P5/32
- [-1, 3, C3, [1024]],
- [-1, 3, SOCA, [1024]],
- [-1, 1, SPPF, [1024, 5]], # 13
- ]
- # YOLOv5 v6.0 head
- head:
- [[-1, 1, Conv, [512, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 9], 1, Concat, [1]], # cat backbone P4
- [-1, 3, C3, [512, False]], # 17
- [-1, 1, Conv, [256, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 6], 1, Concat, [1]], # cat backbone P3
- [-1, 3, C3, [256, False]], # 21 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]],
- [[-1, 18], 1, Concat, [1]], # cat head P4
- [-1, 3, C3, [512, False]], # 24 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]],
- [[-1, 14], 1, Concat, [1]], # cat head P5
- [-1, 3, C3, [1024, False]], # 27 (P5/32-large)
- [[21, 24, 27], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
第④步:验证是否加入成功
同样的方法,我们来运行一下yolo.py

OK~收工!