CBAM代注意力机制代码
时间: 2023-08-14 13:13:23 浏览: 166
以下是一个简单的 CBAM(Convolutional Block Attention Module)注意力机制的代码示例:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class ChannelAttention(nn.Module):
def __init__(self, in_channels, reduction_ratio=16):
super(ChannelAttention, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(
nn.Linear(in_channels, in_channels // reduction_ratio),
nn.ReLU(),
nn.Linear(in_channels // reduction_ratio, in_channels)
)
def forward(self, x):
avg_out = self.fc(self.avg_pool(x).view(x.size(0), -1))
max_out = self.fc(self.max_pool(x).view(x.size(0), -1))
out = avg_out + max_out
out = torch.sigmoid(out).unsqueeze(2).unsqueeze(3).expand_as(x)
return out * x
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super(SpatialAttention, self).__init__()
padding = kernel_size // 2
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
def forward(self, x):
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
out = torch.cat([avg_out, max_out], dim=1)
out = self.conv(out)
out = torch.sigmoid(out)
return out * x
class CBAM(nn.Module):
def __init__(self, in_channels, reduction_ratio=16, kernel_size=7):
super(CBAM, self).__init__()
self.channel_attention = ChannelAttention(in_channels, reduction_ratio)
self.spatial_attention = SpatialAttention(kernel_size)
def forward(self, x):
out = self.channel_attention(x)
out = self.spatial_attention(out)
return out
# 使用示例
input_tensor = torch.randn(1, 64, 32, 32) # 输入数据的尺寸为 [batch_size, channels, height, width]
cbam = CBAM(64) # 创建 CBAM 注意力模块实例
output_tensor = cbam(input_tensor) # 输入数据通过 CBAM 注意力模块获取输出
print(output_tensor.size()) # 打印输出数据的尺寸
```
这段代码实现了 CBAM 注意力机制,包括通道注意力和空间注意力。通道注意力模块用于对输入张量的不同通道进行自适应的加权融合,空间注意力模块用于对输入张量的空间维度进行自适应的加权融合。最后,通过将通道注意力和空间注意力的输出相乘,实现了综合的 CBAM 注意力机制。
阅读全文
相关推荐
![-](https://img-home.csdnimg.cn/images/20241231044955.png)
![-](https://img-home.csdnimg.cn/images/20241231045053.png)
![-](https://img-home.csdnimg.cn/images/20241231044930.png)
![zip](https://img-home.csdnimg.cn/images/20241231045053.png)
![zip](https://img-home.csdnimg.cn/images/20241231045053.png)
![7z](https://img-home.csdnimg.cn/images/20241231044736.png)
![zip](https://img-home.csdnimg.cn/images/20241231045053.png)
![-](https://img-home.csdnimg.cn/images/20241231044955.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![-](https://img-home.csdnimg.cn/images/20241226111658.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)