import numpy as np
import torch
from torch import nn
from torch.nn import init
# CBAM Attention
# 方法出处 2018 ECCV 《CBAM: Convolutional Block Attention Module》
# 通道注意力
class ChannelAttention(nn.Module):
def __init__(self, channel, reduction=16):
# 所有继承于nn.Module的模型都要写这句话
super(ChannelAttention, self).__init__()
# 这个AdaptiveAvgPool2d会将输入特征图的宽和高
# 自动的池化到我们在AdaptiveAvgPool2d参数中指定的大小
# 比如
# m = nn.AdaptiveAvgPool2d((5, 7))
# input = torch.randn(1, 64, 8, 9)
# output = m(input)
# print(output.size())
# 最后会输出[1,64,5,7]
# 如果指定的最后宽和高相等可以只写一个
# 比如
# nn.AdaptiveAvgPool2d((1,1))==nn.AdaptiveAvgPool2d(1)
self.maxpool = nn.AdaptiveMaxPool2d(1)
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.se = nn.Sequential(
# 论文中说这里用的是线性层,不是卷积层
# 但是1*1的卷积层相当于全连接
nn.Conv2d(channel, channel // reduction, 1, bias=False),
nn.ReLU(),
nn.Conv2d(channel // reduction, channel, 1, bias=False)
)
# 激活
self.sigmoid = nn.Sigmoid()
def forward(self, x):
max_result = self.maxpool(x)
avg_result = self.avgpool(x)
max_out = self.se(max_result)
avg_out = self.se(avg_result)
output = self.sigmoid(max_out + avg_out)
return output
# 空间注意力
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super(SpatialAttention, self).__init__()
# 输入通道数是2,max_pool生成一个,avg_pool生成一个
# 输出1个空间注意力特征
# 卷积核大小是指定的kernel_size
# 进行了填充
self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
# 激活
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# 按照论文要求进行最大池化
# 表示沿通道求最大值
max_result, _ = torch.max(x, dim=1, keepdim=True)
# 平均池化
# 同理
avg_result = torch.mean(x, dim=1, keepdim=True)
result = torch.cat([max_result, avg_result], 1)
# 获取空间注意力特征
output = self.conv(result)
output = self.sigmoid(output)
return output
class CBAMBlock(nn.Module):
def __init__(self, channel=512, reduction=16, kernel_size=49):
super(CBAMBlock, self).__init__()
# 组合通道注意力和空间注意力
self.ca = ChannelAttention(channel=channel, reduction=reduction)
self.sa = SpatialAttention(kernel_size=kernel_size)
# 初始化层
def init_weights(self):
# 遍历当前模型的每一层
for m in self.modules():
# 如果是卷积层
if isinstance(m, nn.Conv2d):
# kaiming初始化
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
# 前向传递建立计算图
def forward(self, x):
# 获取输入特诊图的批大小,通道数
b, c, _, _ = x.size()
# 最后有一个残差连接
residual = x
# 计算通道注意力
out = x * self.ca(x)
# 计算空间注意力
out = out * self.sa(out)
# 残差连接
return out + residual
if __name__ == '__main__':
# 可以将这个输入当做卷积层输出的特征图
# 维度[50,512,7,7]
# 50代表批大小
# 512代表通道数
# 7*7是宽高
input = torch.randn(50, 512, 7, 7)
# 获取通道数
kernel_size = input.shape[2]
# 进行通道注意力,空间注意力
cbam = CBAMBlock(channel=512, reduction=16, kernel_size=kernel_size)
output = cbam(input)
print(output.shape)
2018 ECCV 《CBAM: Convolutional Block Attention Module》 PyTorch实现
最新推荐文章于 2024-08-14 09:12:50 发布