序列到序列模型与注意力机制

|
class Seq2SeqModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        self.encoder_embedding = nn.Embedding(src_vocab_size, embedding_dim)
        self.encoder_lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, embedding_dim)
        self.decoder_lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.output_layer = nn.Linear(hidden_dim, tgt_vocab_size)
    
    def forward(self, src, tgt):
        # 编码
        src_embedded = self.encoder_embedding(src)
        _, (h_n, c_n) = self.encoder_lstm(src_embedded)
        
        # 解码
        tgt_embedded = self.decoder_embedding(tgt)
        decoder_output, _ = self.decoder_lstm(tgt_embedded, (h_n, c_n))
        
        # 输出层
        output = self.output_layer(decoder_output)
        return output

|
class AttentionDecoder(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, output_dim):
        super().__init__()
        self.embedding = nn.Embedding(output_dim, embedding_dim)
        self.lstm = nn.LSTMCell(embedding_dim + hidden_dim, hidden_dim)
        self.attention = nn.Linear(hidden_dim * 2, 1)
        self.output_layer = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, prev_word, prev_hidden, encoder_outputs):
        """
        prev_word: 前一个词，shape (batch,)
        prev_hidden: 前一时刻隐藏状态，shape (batch, hidden_dim)
        encoder_outputs: 所有编码器隐藏状态，shape (batch, src_len, hidden_dim)
        """
        batch_size = encoder_outputs.shape[0]
        src_len = encoder_outputs.shape[1]
        
        # 计算注意力权重
        # 将prev_hidden扩展为(batch, src_len, hidden_dim)
        prev_hidden_expanded = prev_hidden.unsqueeze(1).expand(-1, src_len, -1)
        
        # 拼接并通过注意力网络
        combined = torch.cat([encoder_outputs, prev_hidden_expanded], dim=2)
        attention_scores = self.attention(combined).squeeze(2)  # (batch, src_len)
        attention_weights = F.softmax(attention_scores, dim=1)  # (batch, src_len)
        
        # 计算上下文向量
        context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1)  # (batch, hidden_dim)
        
        # 嵌入当前输入词
        embedded = self.embedding(prev_word)  # (batch, emb_dim)
        
        # LSTM输入是[embedded, context]拼接
        lstm_input = torch.cat([embedded, context], dim=1)
        hidden_new, cell_new = self.lstm(lstm_input, (prev_hidden, prev_cell))
        
        # 输出
        output = self.output_layer(hidden_new)
        
        return output, hidden_new, cell_new, attention_weights

|
def beam_search_decode(model, src, beam_width=3, max_length=50):
    """
    集束搜索解码
    """
    # 编码
    context = model.encode(src)
    
    # 初始化：起始标记
    sequences = [[BOS_token]]
    scores = [0.0]
    
    for _ in range(max_length):
        all_candidates = []
        
        # 扩展每个候选序列
        for i in range(len(sequences)):
            seq = sequences[i]
            score = scores[i]
            
            if seq[-1] == EOS_token:  # 已结束的序列
                all_candidates.append((score, seq))
                continue
            
            # 预测下一个词的概率
            probs = model.decode_step(seq, context)
            
            # 取top-K个词
            top_k_probs, top_k_words = torch.topk(probs, beam_width)
            
            for j in range(beam_width):
                candidate_seq = seq + [top_k_words[j].item()]
                candidate_score = score + np.log(top_k_probs[j].item())
                all_candidates.append((candidate_score, candidate_seq))
        
        # 按分数排序，保留top-B
        ordered = sorted(all_candidates, key=lambda x: x[0], reverse=True)
        sequences = [seq for _, seq in ordered[:beam_width]]
        scores = [score for score, _ in ordered[:beam_width]]
        
        # 如果所有序列都结束了，提前停止
        if all(seq[-1] == EOS_token for seq in sequences):
            break
    
    # 返回分数最高的序列
    return sequences[0]

|
# Transformer的自注意力（简化版）
def self_attention(Q, K, V):
    """
    Q, K, V: (batch, seq_len, d_model)
    """
    d_k = Q.shape[-1]
    
    # 计算注意力分数
    scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(d_k)
    
    # Softmax得到注意力权重
    attention_weights = F.softmax(scores, dim=-1)
    
    # 加权和
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights