| 201 | |
| 202 | |
| 203 | class ChannelAttentionModule(nn.Module): |
| 204 | def __init__(self, channel, ratio=16): |
| 205 | super(ChannelAttentionModule, self).__init__() |
| 206 | self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| 207 | self.max_pool = nn.AdaptiveMaxPool2d(1) |
| 208 | |
| 209 | self.shared_MLP = nn.Sequential( |
| 210 | nn.Conv2d(channel, channel // ratio, 1, bias=False), |
| 211 | nn.ReLU(), |
| 212 | nn.Conv2d(channel // ratio, channel, 1, bias=False) |
| 213 | ) |
| 214 | self.sigmoid = nn.Sigmoid() |
| 215 | |
| 216 | def forward(self, x): |
| 217 | avgout = self.shared_MLP(self.avg_pool(x)) |
| 218 | maxout = self.shared_MLP(self.max_pool(x)) |
| 219 | return self.sigmoid(avgout + maxout) |
| 220 | |
| 221 | |
| 222 | class SpatialAttentionModule(nn.Module): |