(self, query_dim, key_dim, value_dim)
| 231 | |
| 232 | class AttLayer(nn.Module): |
| 233 | def __init__(self, query_dim, key_dim, value_dim): |
| 234 | super(AttLayer, self).__init__() |
| 235 | self.W_q = nn.Linear(query_dim, value_dim) |
| 236 | self.W_k = nn.Linear(key_dim, value_dim, bias=False) |
| 237 | self.W_v = nn.Linear(key_dim, value_dim) |
| 238 | |
| 239 | self.softmax = nn.Softmax(dim=1) |
| 240 | self.dim = value_dim |
| 241 | |
| 242 | self.W_q.apply(init_weight) |
| 243 | self.W_k.apply(init_weight) |
| 244 | self.W_v.apply(init_weight) |
| 245 | |
| 246 | def forward(self, query, key_mat): |
| 247 | ''' |