(self, x, proxy)
| 93 | ) |
| 94 | |
| 95 | def forward(self, x, proxy): |
| 96 | batch_size, h, w = x.size(0), x.size(2), x.size(3) |
| 97 | if self.scale > 1: |
| 98 | x = self.pool(x) |
| 99 | |
| 100 | query = self.f_pixel(x).view(batch_size, self.key_channels, -1) |
| 101 | query = query.permute(0, 2, 1) |
| 102 | key = self.f_object(proxy).view(batch_size, self.key_channels, -1) |
| 103 | value = self.f_down(proxy).view(batch_size, self.key_channels, -1) |
| 104 | value = value.permute(0, 2, 1) |
| 105 | |
| 106 | sim_map = torch.matmul(query, key) |
| 107 | sim_map = (self.key_channels**-.5) * sim_map |
| 108 | sim_map = F.softmax(sim_map, dim=-1) |
| 109 | |
| 110 | # add bg context ... |
| 111 | context = torch.matmul(sim_map, value) |
| 112 | context = context.permute(0, 2, 1).contiguous() |
| 113 | context = context.view(batch_size, self.key_channels, *x.size()[2:]) |
| 114 | context = self.f_up(context) |
| 115 | if self.scale > 1: |
| 116 | context = F.interpolate(input=context, size=(h, w), mode='bilinear', |
| 117 | align_corners=cfg.MODEL.ALIGN_CORNERS) |
| 118 | |
| 119 | return context |
| 120 | |
| 121 | |
| 122 | class SpatialOCR_Module(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected