| 152 | return x.flatten(self.start_dim) |
| 153 | |
| 154 | class SqueezeExcitation(nn.Module): |
| 155 | def __init__(self, channels, reduction=16): |
| 156 | super().__init__() |
| 157 | self.fc1 = nn.Linear(channels, channels // reduction) |
| 158 | self.fc2 = nn.Linear(channels // reduction, channels) |
| 159 | |
| 160 | def forward(self, x): |
| 161 | b, c, _, _ = x.size() |
| 162 | se = x.mean([2, 3]) |
| 163 | se = F.relu(self.fc1(se)) |
| 164 | se = torch.sigmoid(self.fc2(se)) |
| 165 | return x * se.view(b, c, 1, 1) |
| 166 | |
| 167 | class SpatialAttention(nn.Module): |
| 168 | def __init__(self, kernel_size=7): |