PyTorch 代码注释:融合层和注意力机制实现
定义一个求和融合层的类
class SumFusion(nn.Module): def init(self, dim): super(SumFusion, self).init() # 初始化层参数 # dim: 特征图的通道数
# 定义前向传播函数
def forward(self, in_feats):
# 对输入特征图列表进行求和操作
return sum(in_feats)
定义一个拼接融合层的类
class CatFusion(nn.Module): def init(self, dim, height=2): super(CatFusion, self).init() # 初始化层参数 self.height = height # 输入特征图的数量 # 定义一个卷积层,将拼接后的特征图进行降维 self.conv = nn.Conv2d(dim * height, dim, kernel_size=1)
# 定义前向传播函数
def forward(self, in_feats):
# 将输入特征图列表进行拼接操作
in_feats = torch.cat(in_feats, dim=1)
# 使用卷积层对拼接后的特征图进行降维
out = self.conv(in_feats)
# 返回降维后的特征图
return out
定义一个通道注意力机制模块
class CALayer(nn.Module): def init(self, channel, reduction=8): super(CALayer, self).init() # 初始化层参数 d = max(channel // reduction, 4) # 通道注意力机制的中间层通道数
# 定义一个自适应平均池化层,将特征图降到1x1大小
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# 定义一个通道注意力机制的结构,包括两个卷积层和一个sigmoid函数
self.ca = nn.Sequential(
# 第一个卷积层,将通道数降到d
nn.Conv2d(channel, d, 1, bias=False),
# ReLU激活函数
nn.ReLU(inplace=True),
# 第二个卷积层,将通道数恢复到channel
nn.Conv2d(d, channel, 1, bias=False),
# Sigmoid函数,将输出归一化到0~1之间
nn.Sigmoid()
)
# 定义前向传播函数
def forward(self, x):
# 使用自适应平均池化层对输入特征图进行全局平均池化
y = self.avg_pool(x)
# 使用通道注意力机制模块对池化后的特征图进行处理
y = self.ca(y)
# 将通道注意力机制的输出与原特征图相乘
return x * y
定义一个空间注意力机制模块
class ECALayer(nn.Module): def init(self, channel, k_size=3): super(ECALayer, self).init() # 初始化层参数 # channel: 特征图的通道数 # k_size: 卷积核大小
# 定义一个自适应平均池化层,将特征图降到1x1大小
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# 定义一个一维卷积层,用于计算每个通道的权重
self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
# 定义一个sigmoid函数,将权重归一化到0~1之间
self.sigmoid = nn.Sigmoid()
# 定义前向传播函数
def forward(self, x):
# 使用自适应平均池化层对输入特征图进行全局平均池化
y = self.avg_pool(x)
# 将y的形状从 [batch_size, channel, 1, 1] 转换为 [batch_size, channel, 1]
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
# 使用sigmoid函数对卷积输出进行归一化
y = self.sigmoid(y)
# 将权重扩展到每个像素点,并与原特征图相乘,得到加强后的特征图
return x * y.expand_as(x)
原文地址: https://www.cveoy.top/t/topic/nw37 著作权归作者所有。请勿转载和采集!