generate mask
(self, features, lengths, mask, k)
| 50 | return tags |
| 51 | |
| 52 | def multi_hops(self, features, lengths, mask, k): |
| 53 | '''generate mask''' |
| 54 | max_length = features.shape[1] |
| 55 | mask = mask[:, :max_length] |
| 56 | mask_a = mask.unsqueeze(1).expand([-1, max_length, -1]) |
| 57 | mask_b = mask.unsqueeze(2).expand([-1, -1, max_length]) |
| 58 | mask = mask_a * mask_b |
| 59 | mask = torch.triu(mask).unsqueeze(3).expand([-1, -1, -1, self.args.class_num]) |
| 60 | |
| 61 | '''save all logits''' |
| 62 | logits_list = [] |
| 63 | logits = self._cls_logits(features) |
| 64 | logits_list.append(logits) |
| 65 | |
| 66 | for i in range(k): |
| 67 | #probs = torch.softmax(logits, dim=3) |
| 68 | probs = logits |
| 69 | logits = probs * mask |
| 70 | |
| 71 | logits_a = torch.max(logits, dim=1)[0] |
| 72 | logits_b = torch.max(logits, dim=2)[0] |
| 73 | logits = torch.cat([logits_a.unsqueeze(3), logits_b.unsqueeze(3)], dim=3) |
| 74 | logits = torch.max(logits, dim=3)[0] |
| 75 | |
| 76 | logits = logits.unsqueeze(2).expand([-1,-1, max_length, -1]) |
| 77 | logits_T = logits.transpose(1, 2) |
| 78 | logits = torch.cat([logits, logits_T], dim=3) |
| 79 | |
| 80 | new_features = torch.cat([features, logits, probs], dim=3) |
| 81 | features = self.feature_linear(new_features) |
| 82 | logits = self._cls_logits(features) |
| 83 | logits_list.append(logits) |
| 84 | return logits_list |
| 85 | |
| 86 | def forward(self, sentence_tokens, lengths, mask): |
| 87 | embedding = self._get_embedding(sentence_tokens, mask) |