pytorch——在lstm模型中加入自注意力机制

首先先说为什么要加入自注意力机制。
自注意力机制是将输入的内容进行查询，并根据每一个词的相关性赋予不同的权重。而后在计算中会围绕每个词的权重进行计算。某种意义上说，他也是将长句变短的一个方法，所以在计算效率上，比单纯使用lstm要节省资源。

自注意力层是加载lstm和Linear层中间的一层，用于对lstm层输出的结果进行权重赋予，再将最后结果传入最后的Linear层进行计算。

class SelfAttention(nn.Module):
    def __init__(self, hidden_dim):
        super(SelfAttention, self).__init__()
        self.projection = nn.Sequential(
            nn.Linear(hidden_dim, 64),   # hidden_dim 为lstm的隐藏层数
            nn.ReLU(True),
            nn.Linear(64, 1)
        )

    def forward(self, encoder_outputs):
        # encoder_outputs: (batch_size, sequence_length, hidden_dim)
        # 计算注意力得分，赋予权重
        energy = self.projection(encoder_outputs)  # (batch_size, sequence_length, 1)
        weights = F.softmax(energy.squeeze(-1), dim=1)  # (batch_size, sequence_length)

        # 应用注意力权重
        outputs = (encoder_outputs * weights.unsqueeze(-1)).sum(dim=1)  # (batch_size, hidden_dim)
        return outputs, weights

然后再将这个自注意力model对LSTM的输出结果进行计算。

class Attention_LSTM_Model(nn.Module):
    def __init__(self, config):
        super(Attention_LSTM_Model, self).__init__()
       
        # LSTM网络
        # config.embed：词向量的输出长度=300，它是LSTM的输入
        # config.hidden_size：隐藏层输出特征的长度
        # config.num_layers：隐藏层的层数
        # bidirectional：双向网络，这里面没有启用，也不建议启用
        # batch_first： [batch_size, seq_len, embeding] 如果这个值为False，则输出[seq_len,batch_size,embeding]
        # dropout：随机丢弃
        self.lstm = nn.LSTM(config.embed, config.hidden_size, config.num_layers,
                             batch_first=True, dropout=config.dropout)
        #添加自注意力层和注意力层
        self.self_attention = SelfAttention(config.hidden_size)
        self.attention = Attention(config.hidden_size)
        # 全连接分类网络
        # 使用隐层当前时刻的输出作为全连接的输入。
        # config.hidden_size * 2：双向LSTM的输出是隐层特征输出的2倍
        self.fc = nn.Linear(config.hidden_size, config.num_classes)

    def forward(self, x):
        lstm_out, (hidden, _) = self.lstm()  # 输出的shape为[64, 11, 256]  64是batchsize，11是11个文字，256是隐藏层
        # lstm_out: (batch_size, sequence_length, hidden_dim)
        # 自注意力
        self_att_out, self_weights = self.self_attention(lstm_out)
        
        # 通过全连接层
        output = self.fc(self_att_out)
        return output

同时也可以给lstm模型添加其他注意力机制，最后再和这个att_out相加就可以了。