(self, feats, probs)
| 32 | self.scale = scale |
| 33 | |
| 34 | def forward(self, feats, probs): |
| 35 | batch_size, c, _, _ = probs.size(0), probs.size(1), probs.size(2), \ |
| 36 | probs.size(3) |
| 37 | |
| 38 | # each class image now a vector |
| 39 | probs = probs.view(batch_size, c, -1) |
| 40 | feats = feats.view(batch_size, feats.size(1), -1) |
| 41 | |
| 42 | feats = feats.permute(0, 2, 1) # batch x hw x c |
| 43 | probs = F.softmax(self.scale * probs, dim=2) # batch x k x hw |
| 44 | ocr_context = torch.matmul(probs, feats) |
| 45 | ocr_context = ocr_context.permute(0, 2, 1).unsqueeze(3) |
| 46 | return ocr_context |
| 47 | |
| 48 | |
| 49 | class ObjectAttentionBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected