Global Context Extractor for CGNet. This class is employed to refine the joint feature of both local feature and surrounding context. Args: channel (int): Number of input feature channels. reduction (int): Reductions for global context extractor. Default: 16. wi
| 12 | |
| 13 | |
| 14 | class GlobalContextExtractor(nn.Module): |
| 15 | """Global Context Extractor for CGNet. |
| 16 | |
| 17 | This class is employed to refine the joint feature of both local feature |
| 18 | and surrounding context. |
| 19 | |
| 20 | Args: |
| 21 | channel (int): Number of input feature channels. |
| 22 | reduction (int): Reductions for global context extractor. Default: 16. |
| 23 | with_cp (bool): Use checkpoint or not. Using checkpoint will save some |
| 24 | memory while slowing down the training speed. Default: False. |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, channel, reduction=16, with_cp=False): |
| 28 | super(GlobalContextExtractor, self).__init__() |
| 29 | self.channel = channel |
| 30 | self.reduction = reduction |
| 31 | assert reduction >= 1 and channel >= reduction |
| 32 | self.with_cp = with_cp |
| 33 | self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| 34 | self.fc = nn.Sequential( |
| 35 | nn.Linear(channel, channel // reduction), nn.ReLU(inplace=True), |
| 36 | nn.Linear(channel // reduction, channel), nn.Sigmoid()) |
| 37 | |
| 38 | def forward(self, x): |
| 39 | |
| 40 | def _inner_forward(x): |
| 41 | num_batch, num_channel = x.size()[:2] |
| 42 | y = self.avg_pool(x).view(num_batch, num_channel) |
| 43 | y = self.fc(y).view(num_batch, num_channel, 1, 1) |
| 44 | return x * y |
| 45 | |
| 46 | if self.with_cp and x.requires_grad: |
| 47 | out = cp.checkpoint(_inner_forward, x) |
| 48 | else: |
| 49 | out = _inner_forward(x) |
| 50 | |
| 51 | return out |
| 52 | |
| 53 | |
| 54 | class ContextGuidedBlock(nn.Module): |