基于Python的自然语言处理系列（27）：LSTM与Attention机制从零开始实现

        在这一篇中，我们将深入探讨LSTM及其变体的结构，并从零开始实现这些模型。为了帮助大家更好地理解，我们会逐步实现Vanilla LSTM、Coupled Gate LSTM、Peephole LSTM以及双向LSTM（BiLSTM）。此外，还会实现Attention机制，包括General Attention和Self-Attention，帮助我们理解它们在自然语言处理中的应用。

1. LSTM及其变体

1.1 LSTM工作机制

        LSTM的主要作用是通过“记忆单元”（cell state）来保存信息，并通过三个“门控机制”——遗忘门、输入门和输出门，来控制哪些信息应该保留、更新或输出。

• 遗忘门：决定从上一时间步的记忆中舍弃哪些信息。
• 输入门：决定当前时间步哪些新的信息需要写入到记忆中。
• 输出门：根据更新后的记忆状态，决定输出什么信息。

        通过这些机制，LSTM可以很好地捕捉序列中的长期依赖关系，避免信息随着时间步的增加而丢失。

1.2 代码实现LSTM

        下面是一个从零开始实现的LSTM单元。

import torch
import torch.nn as nn
import math

class LSTMCell(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(LSTMCell, self).__init__()
        self.hidden_dim = hidden_dim
        
        # 初始化权重矩阵
        self.W_f = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
        self.U_f = nn.Parameter(torch.Tensor(input_dim, hidden_dim))
        self.b_f = nn.Parameter(torch.Tensor(hidden_dim))
        
        self.W_i = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
        self.U_i = nn.Parameter(torch.Tensor(input_dim, hidden_dim))
        self.b_i = nn.Parameter(torch.Tensor(hidden_dim))
        
        self.W_o = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
        self.U_o = nn.Parameter(torch.Tensor(input_dim, hidden_dim))
        self.b_o = nn.Parameter(torch.Tensor(hidden_dim))
        
        self.W_g = nn.Parameter(torch.Tensor(hidden_dim, hidden_dim))
        self.U_g = nn.Parameter(torch.Tensor(input_dim, hidden_dim))
        self.b_g = nn.Parameter(torch.Tensor(hidden_dim))
        
        self.init_weights()

    def init_weights(self):
        stdv = 1.0 / math.sqrt(self.hidden_dim)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def forward(self, x, init_states=None):
        bs, seq_len, _ = x.shape
        outputs = []
        
        if init_states is None:
            h_t = torch.zeros(bs, self.hidden_dim).to(x.device)
            c_t = torch.zeros(bs, self.hidden_dim).to(x.device)
        else:
            h_t, c_t = init_states
        
        for t in range(seq_len):
            x_t = x[:, t, :]
            f_t = torch.sigmoid(h_t @ self.W_f + x_t @ self.U_f + self.b_f)
            i_t = torch.sigmoid(h_t @ self.W_i + x_t @ self.U_i + self.b_i)
            o_t = torch.sigmoid(h_t @ self.W_o + x_t @ self.U_o + self.b_o)
            g_t = torch.tanh(h_t @ self.W_g + x_t @