| 72 | |
| 73 | |
| 74 | class SEBlock(nn.Module): |
| 75 | def __init__(self, ch_in, ch_out): |
| 76 | super().__init__() |
| 77 | |
| 78 | self.main = nn.Sequential( nn.AdaptiveAvgPool2d(4), |
| 79 | conv2d(ch_in, ch_out, 4, 1, 0, bias=False), Swish(), |
| 80 | conv2d(ch_out, ch_out, 1, 1, 0, bias=False), nn.Sigmoid() ) |
| 81 | |
| 82 | def forward(self, feat_small, feat_big): |
| 83 | return feat_big * self.main(feat_small) |
| 84 | |
| 85 | |
| 86 | class InitLayer(nn.Module): |