import torch
import torch.nn as nn
import torch.nn.functional as F
class ChannelAttentionModule(nn.Module):
def __init__(self, in_channels, reduction_ratio=16):
super(ChannelAttentionModule, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(in_channels, in_channels // reduction_ratio),
nn.ReLU(inplace=True),
nn.Linear(in_channels // reduction_ratio, in_channels)
)
def forward(self, x):
avg_pool = self.avg_pool(x).view(x.size(0), -1)
channel_att = torch.sigmoid(self.fc(avg_pool)).view(x.size(0), x.size(1), 1, 1)
return x * channel_att
class SpatialAt
CBAM注意力机制的示例
最新推荐文章于 2024-08-24 14:57:50 发布
![](https://img-home.csdnimg.cn/images/20240711042549.png)