一种输入[batch, seq_len1, hidden_dim]输出[batch, seq_len2, hidden_dim]的self-attention的pytorch实现

一种输入[batch, seq_len1, hidden_dim]输出[batch, seq_len2, hidden_dim]的self-attention的pytorch实现

class Attention(nn.Module): """ inputs是[batch, seq_len1, hidden_dim] labels_num是seq_len2 """ def __init__(self, labels_num, hidden_size): super(Attention, self).__init__() self.attention = nn.Linear(hidden_size, labels_num, bias=False) nn.init.xavier_uniform_(self.attention.weight) def forward(self, inputs, masks): masks = torch.unsqueeze(masks, 1) # [batch, 1, seq_len1] attention = self.attention(inputs).transpose(1, 2).masked_fill(1.0 - masks, -np.inf) # attention 是 [batch, labels_num, seq_len1] attention = F.softmax(attention, -1) return attention @ inputs # return结果 [batch, labels_num, hidden_size]

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。