Squeeze-and-Excitation (SE) block w/ Swish: AvgPool, FC, Swish, FC, Sigmoid.
| 34 | |
| 35 | |
| 36 | class SE(nn.Module): |
| 37 | """Squeeze-and-Excitation (SE) block w/ Swish: AvgPool, FC, Swish, FC, Sigmoid.""" |
| 38 | |
| 39 | def _round_width(self, width, multiplier, min_width=8, divisor=8): |
| 40 | """ |
| 41 | Round width of filters based on width multiplier |
| 42 | Args: |
| 43 | width (int): the channel dimensions of the input. |
| 44 | multiplier (float): the multiplication factor. |
| 45 | min_width (int): the minimum width after multiplication. |
| 46 | divisor (int): the new width should be dividable by divisor. |
| 47 | """ |
| 48 | if not multiplier: |
| 49 | return width |
| 50 | |
| 51 | width *= multiplier |
| 52 | min_width = min_width or divisor |
| 53 | width_out = max( |
| 54 | min_width, int(width + divisor / 2) // divisor * divisor |
| 55 | ) |
| 56 | if width_out < 0.9 * width: |
| 57 | width_out += divisor |
| 58 | return int(width_out) |
| 59 | |
| 60 | def __init__(self, dim_in, ratio, relu_act=True): |
| 61 | """ |
| 62 | Args: |
| 63 | dim_in (int): the channel dimensions of the input. |
| 64 | ratio (float): the channel reduction ratio for squeeze. |
| 65 | relu_act (bool): whether to use ReLU activation instead |
| 66 | of Swish (default). |
| 67 | divisor (int): the new width should be dividable by divisor. |
| 68 | """ |
| 69 | super(SE, self).__init__() |
| 70 | self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1)) |
| 71 | dim_fc = self._round_width(dim_in, ratio) |
| 72 | self.fc1 = nn.Conv3d(dim_in, dim_fc, 1, bias=True) |
| 73 | self.fc1_act = nn.ReLU() if relu_act else Swish() |
| 74 | self.fc2 = nn.Conv3d(dim_fc, dim_in, 1, bias=True) |
| 75 | |
| 76 | self.fc2_sig = nn.Sigmoid() |
| 77 | |
| 78 | def forward(self, x): |
| 79 | x_in = x |
| 80 | for module in self.children(): |
| 81 | x = module(x) |
| 82 | return x_in * x |