LSTM多头注意力机制pytorch

### 实现 LSTM 与多头注意力机制结合

为了在 PyTorch 中实现 LSTM 和多头注意力机制的结合，可以按照以下方式设计模型架构。这种组合能够充分利用 LSTM 的时间序列处理能力和多头注意力机制的信息捕捉能力。

#### 构建 LSTM 层
首先定义一个标准的 LSTM 层用于处理输入的时间序列数据：

import torch.nn as nn

class LSTMLayer(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=1):
        super(LSTMLayer, self).__init__()
        self.lstm = nn.LSTM(input_size=input_size,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            batch_first=True)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        return lstm_out

此部分负责接收并编码输入序列到固定长度表示形式[^1]。

#### 添加多头自注意层
接着引入基于 `nn.MultiheadAttention` 组件构建的多头注意力模块，该组件允许不同位置之间的交互更加灵活有效:

from typing import Tuple

class MultiHeadSelfAttentionLayer(nn.Module):
    def __init__(self, embed_dim: int, n_heads: int):
        super(MultiHeadSelfAttentionLayer, self).__init__()
        self.multi_attn = nn.MultiheadAttention(embed_dim=embed_dim,
                                                num_heads=n_heads,
                                                batch_first=True)

    def forward(self, query_key_value: torch.Tensor) -> Tuple[torch.Tensor]:
        attn_output, _ = self.multi_attn(query=query_key_value,
                                         key=query_key_value,
                                         value=query_key_value)
        return attn_output

这里实现了多头自我关注功能，其中每个头部独立工作于不同的子空间内提取特征[^3]。

#### 整合两者的混合模型
最后一步就是把上述两个主要组成部分结合起来形成完整的网络框架：

class LSTMWithMultiHeadAttnModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim, lstm_hidden_dim, attn_head_num):
        super().__init__()

        # Embedding layer to convert token indices into dense vectors of fixed size.
        self.embedding = nn.Embedding(vocab_size, embedding_dim)

        # Define the LSTM and Attention layers based on previous definitions.
        self.lstm_layer = LSTMLayer(input_size=embedding_dim,
                                    hidden_size=lstm_hidden_dim)
        
        self.attn_layer = MultiHeadSelfAttentionLayer(
            embed_dim=lstm_hidden_dim, 
            n_heads=attn_head_num
        )

        # Output layer (for simplicity assuming binary classification task).
        self.fc_out = nn.Linear(lstm_hidden_dim, 1)

    def forward(self, text_input_ids):

        embedded_seq = self.embedding(text_input_ids)
        lstm_encoded = self.lstm_layer(embedded_seq)
        attended_features = self.attn_layer(lstm_encoded)
        logits = self.fc_out(attended_features[:, -1, :])  # Use last time step output
        
        return logits.squeeze(-1)

这段代码展示了如何创建一个融合了LSTM和多头注意力机制的深度学习分类器，在自然语言处理等领域具有广泛应用前景。