| 58 | |
| 59 | |
| 60 | class SelfAttention(torch.nn.Module): |
| 61 | def __init__(self, args): |
| 62 | super(SelfAttention,self).__init__() |
| 63 | self.args = args |
| 64 | self.linear_q = torch.nn.Linear(args.lstm_dim * 2, args.lstm_dim * 2) |
| 65 | # self.linear_k = torch.nn.Linear(configs.BILSTM_DIM * 2, configs.BILSTM_DIM * 2) |
| 66 | # self.linear_v = torch.nn.Linear(configs.BILSTM_DIM * 2, configs.BILSTM_DIM * 2) |
| 67 | # self.w_query = torch.nn.Linear(configs.BILSTM_DIM * 2, 50) |
| 68 | # self.w_value = torch.nn.Linear(configs.BILSTM_DIM * 2, 50) |
| 69 | self.w_query = torch.nn.Linear(args.cnn_dim, 50) |
| 70 | self.w_value = torch.nn.Linear(args.cnn_dim, 50) |
| 71 | self.v = torch.nn.Linear(50, 1, bias=False) |
| 72 | |
| 73 | def forward(self, query, value, mask): |
| 74 | # attention_states = self.linear_q(query) |
| 75 | # attention_states_T = self.linear_k(values) |
| 76 | attention_states = query |
| 77 | attention_states_T = value |
| 78 | attention_states_T = attention_states_T.permute([0, 2, 1]) |
| 79 | |
| 80 | weights=torch.bmm(attention_states, attention_states_T) |
| 81 | weights = weights.masked_fill(mask.unsqueeze(1).expand_as(weights)==0, float("-inf")) # mask掉每行后面的列 |
| 82 | attention = F.softmax(weights,dim=2) |
| 83 | |
| 84 | # value=self.linear_v(states) |
| 85 | merged=torch.bmm(attention, value) |
| 86 | merged=merged * mask.unsqueeze(2).float().expand_as(merged) |
| 87 | |
| 88 | return merged |
| 89 | |
| 90 | def forward_perceptron(self, query, value, mask): |
| 91 | attention_states = query |
| 92 | attention_states = self.w_query(attention_states) |
| 93 | attention_states = attention_states.unsqueeze(2).expand(-1,-1,attention_states.shape[1], -1) |
| 94 | |
| 95 | attention_states_T = value |
| 96 | attention_states_T = self.w_value(attention_states_T) |
| 97 | attention_states_T = attention_states_T.unsqueeze(2).expand(-1,-1,attention_states_T.shape[1], -1) |
| 98 | attention_states_T = attention_states_T.permute([0, 2, 1, 3]) |
| 99 | |
| 100 | weights = torch.tanh(attention_states+attention_states_T) |
| 101 | weights = self.v(weights).squeeze(3) |
| 102 | weights = weights.masked_fill(mask.unsqueeze(1).expand_as(weights)==0, float("-inf")) # mask掉每行后面的列 |
| 103 | attention = F.softmax(weights,dim=2) |
| 104 | |
| 105 | merged = torch.bmm(attention, value) |
| 106 | merged = merged * mask.unsqueeze(2).float().expand_as(merged) |
| 107 | return merged |