Pytorch注意力机制应用到具体网络方法(闭眼都会版)

文章目录

以YoloV4-tiny为例

解释一下各个部分:

  • 最左边这部分为主干提取网络,功能为特征提取
  • 中间这边部分为FPN,功能是加强特征提取
  • 最后一部分为yolo head,功能为获得我们具体的一个预测结果

需要明白几个点:

  • 注意力机制模块是一个即插即用的模块,理论上是可以添加到任何一个特征图后面
  • 但是,不建议添加到主干部分(即最左边的那部分),主干部分所用的特征是我们后面处理所用的基础,故不建议添加到主干部分
  • 如果添加到主干部分,由于注意力机制模块 它的权值模块是随机初始化的,那主干部分的权值就被破坏了,最开始提取出来的特征就不好用了。
  • 故建议把注意力机制模块添加到主干以外的部分

本节把注意力机制添加到加强网络里面,即上图的中间部分。

添加注意力机制可以添加到上图标注的部分。

要加入的注意力机制代码

这一部分为要加入的注意力机制模块,文件名为attention.py

python 复制代码
import torch
from torch import nn
# 通道注意力机制
class channel_attention(nn.Module):
    def __init__(self,channel,ration=16):   #因为要进行全连接,故需要传入通道数量,及缩放比例
        super(channel_attention,self).__init__()  #初始化
        #定义最大池化层
        self.max_pool = nn.AdaptiveMaxPool2d(1) #输出层的高和宽是1
        #定义平均池化
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        self.fc = nn.Sequential(
            #定义第一次全连接
            nn.Linear(channel,channel // ration ,False),
            nn.ReLU(),
            # 定义第二次全连接
            nn.Linear(channel//ration,channel,False)
        )
        #由于图中的通道注意力机制是连个全连接层相加之后再取sigmoid
        self.sigmoid=nn.Sigmoid()

    #前传部分
    def forward(self,x):
        b,c,h,w=x.size()
        #首先对输入进来的x先进行一个全局最大池化 在进行一个全局平均池化
        max_pool_out=self.max_pool(x).view([b,c])
        avg_pool_out=self.avg_pool(x).view([b,c])
        #然后对两次池化后的结果用共享的全连接层fc进行处理
        max_fc_out=self.fc(max_pool_out)
        avg_fc_out=self.fc(avg_pool_out)
        #最后将上面的两个结果进行相加
        out=max_fc_out + avg_fc_out
        out=self.sigmoid(out).view([b,c,1,1])
        #print(out)
        return out * x
# 空间注意力机制
class spacial_attention(nn.Module):
    def __init__(self,kernel_size=7):   #空间注意力没有通道数,故不用传入channel和ration
        #但是空间注意力会进行一次卷积,故我们需要关注卷积核大小,一般为3或7
        super(spacial_attention,self).__init__()  #初始化
        padding=7//2  #卷积核大小整除输入通道数
        self.conv=nn.Conv2d(2,1,kernel_size,1,padding,bias=False)
        #由图可知输入通道数是2,输出通道数为1,卷积核大小默认设置为7,步长为1,因为不需要压缩特征层阿高和宽

        #由于图中的通道注意力机制是连个全连接层相加之后再取sigmoid
        self.sigmoid=nn.Sigmoid()
    #空间注意力机制前传部分
    def forward(self,x):
        b,c,h,w=x.size()
        max_pool_out,_= torch.max(x,dim=1,keepdim=True)#需要把通道这一维度保留下来,故设置keepdim为True
        #对于pytorch来讲,它的通道是在第一维度,也就是batchsize后面的那个维度故定义dim为1
        mean_pool_out = torch.mean(x,dim = 1,keepdim=True)
        #对最大值和平均值进行一个堆叠
        pool_out = torch.cat([max_pool_out, mean_pool_out],dim=1)
        #对堆叠后的结果取一个卷积
        out=self.conv(pool_out)
        out=self.sigmoid(out)
        print(out)
        return out * x

#把空间注意力机制和通道注意力机制进行一个融合
class Cbam(nn.Module):
    def __init__(self,channel,ratio=16,kernel_size=7):
        super(Cbam,self).__init__()
        #调用已经定义好的2个注意力机制
        self.channel_attention=channel_attention(channel,ratio)
        self.spacial_attention = spacial_attention(kernel_size)
    #融合后机制的前传部分
    def forward(self,x):
        x=self.channel_attention(x)
        x=self.spacial_attention(x)
        return x

在模型文件(yolo.py)中,首行添加如下部分

python 复制代码
from .attention import se_block,cbam_block,eca_block
attention_blocks=[se_block,cbam_block,eca_block]
为何要设置成上面的形式?
为了方便调用,到时候可以直接编写下面的代码调用具体的注意力机制模块
attention_blocks[0]

之后,需要找到yolo.py里面的模型主体部分,大概形式如下代码

python 复制代码
class YoloBody(nn.Module):
	def __init__(self,anchors_mask,num_classes,phi=0)
	#在原来的代码上只是添加了phi,代表我们选用的注意力机制模块,默认情况下为0
		super(YoloBody, self).__init__()
	        self.backbone       = darknet53_tiny(None)
	
	        self.conv_for_P5    = BasicConv(512,256,1)
	        self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
	
	        self.upsample       = Upsample(256,128)
	        self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
	        #下面这部分为自己填写
	     	self.phi    = phi  #这个是自己添加的
	        if 1 <= self.phi and self.phi <= 3:
            self.feat1_att      = attention_block[self.phi - 1](256)  #通道数为256
            self.feat2_att      = attention_block[self.phi - 1](512)#通道数为512
            self.upsample_att   = attention_block[self.phi - 1](128)#通道数为128
            #通道数到底是多少看这个模型的前传部分的通道数为多少
    def forward(self, x):
		#---------------------------------------------------#
		#   生成CSPdarknet53_tiny的主干模型
		#   feat1的shape为26,26,256
		#   feat2的shape为13,13,512
		#---------------------------------------------------#
		feat1, feat2 = self.backbone(x)
		#下面代码为自己填写
		if 1 <= self.phi and self.phi <= 3:#如果满足条件就添加具体的注意力机制
		    feat1 = self.feat1_att(feat1)
		    feat2 = self.feat2_att(feat2)
		#下面代码模型自带
		# 13,13,512 -> 13,13,256
		P5 = self.conv_for_P5(feat2)
		# 13,13,256 -> 13,13,512 -> 13,13,255
		out0 = self.yolo_headP5(P5) 
		
		# 13,13,256 -> 13,13,128 -> 26,26,128
		P5_Upsample = self.upsample(P5)
		# 26,26,256 + 26,26,128 -> 26,26,384
		#上面代码模型自带,下面代码自己编写
		if 1 <= self.phi and self.phi <= 3:
		    P5_Upsample = self.upsample_att(P5_Upsample)
		 #下面代码模型自带
		P4 = torch.cat([P5_Upsample,feat1],axis=1)
		
		# 26,26,384 -> 26,26,256 -> 26,26,255
		out1 = self.yolo_headP4(P4)
		
		return out0, out1

模型中插入注意力机制

相关推荐
AndrewHZ10 分钟前
【图像处理基石】什么是tone mapping?
图像处理·人工智能·算法·计算机视觉·hdr
Ai尚研修-贾莲11 分钟前
基于DeepSeek、ChatGPT支持下的地质灾害风险评估、易发性分析、信息化建库及灾后重建
人工智能·chatgpt
SelectDB技术团队35 分钟前
Apache Doris 2025 Roadmap:构建 GenAI 时代实时高效统一的数据底座
大数据·数据库·数据仓库·人工智能·ai·数据分析·湖仓一体
weixin_4352081635 分钟前
通过 Markdown 改进 RAG 文档处理
人工智能·python·算法·自然语言处理·面试·nlp·aigc
大数据在线39 分钟前
AI重塑云基础设施,亚马逊云科技打造AI定制版IaaS“样板房”
人工智能·云基础设施·ai大模型·亚马逊云科技
hello_ejb339 分钟前
聊聊Spring AI的RetrievalAugmentationAdvisor
人工智能·spring·restful
东方佑1 小时前
利用Python自动化处理PPT样式与结构:从提取到生成
python·自动化·powerpoint
你觉得2051 小时前
浙江大学朱霖潮研究员:《人工智能重塑科学与工程研究》以蛋白质结构预测为例|附PPT下载方法
大数据·人工智能·机器学习·ai·云计算·aigc·powerpoint
橘猫云计算机设计1 小时前
基于springboot的考研成绩查询系统(源码+lw+部署文档+讲解),源码可白嫖!
java·spring boot·后端·python·考研·django·毕业设计
牙牙要健康1 小时前
【目标检测】【深度学习】【Pytorch版本】YOLOV3模型算法详解
pytorch·深度学习·目标检测