SE Module Ref: https://github.com/moskomule/senet.pytorch/blob/master/senet/se_module.py
| 55 | |
| 56 | |
| 57 | class SEModule(nn.Module): |
| 58 | ''' |
| 59 | SE Module |
| 60 | Ref: https://github.com/moskomule/senet.pytorch/blob/master/senet/se_module.py |
| 61 | ''' |
| 62 | |
| 63 | def __init__(self, in_channels_num, reduction_ratio=4): |
| 64 | super(SEModule, self).__init__() |
| 65 | |
| 66 | if in_channels_num % reduction_ratio != 0: |
| 67 | raise ValueError('in_channels_num must be divisible by reduction_ratio(default = 4)') |
| 68 | |
| 69 | self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| 70 | self.fc = nn.Sequential( |
| 71 | nn.Linear(in_channels_num, in_channels_num // reduction_ratio, bias=False), |
| 72 | nn.ReLU(inplace=False), |
| 73 | nn.Linear(in_channels_num // reduction_ratio, in_channels_num, bias=False), |
| 74 | H_sigmoid() |
| 75 | ) |
| 76 | |
| 77 | def forward(self, x): |
| 78 | batch_size, channel_num, _, _ = x.size() |
| 79 | y = self.avg_pool(x).view(batch_size, channel_num) |
| 80 | y = self.fc(y).view(batch_size, channel_num, 1, 1) |
| 81 | return x * y |
| 82 | |
| 83 | |
| 84 | class Bottleneck(nn.Module): |