基于PyTorch官方Transformer翻译教程的Beam Search解码实现咨询

以下代码完全兼容PyTorch官方入门翻译Transformer教程的代码结构，可直接替换教程中原有的贪心解码逻辑，无需修改教程已定义的模型、词典、掩码生成等基础工具。

核心实现代码 #

import torch
import torch.nn as nn
from torch import Tensor

# 假设教程中已定义以下特殊字符索引，如与你的实际定义不符请自行调整
BOS_IDX = 0
EOS_IDX = 1
PAD_IDX = 2

def beam_search_decode(model: nn.Module, src: Tensor, src_mask: Tensor, 
                       max_len: int, beam_size: int = 3, alpha: float = 0.7,
                       device: torch.device = torch.device('cpu')) -> Tensor:
    """
    束搜索解码函数，参数格式与官方教程的greedy_decode函数完全对齐
    Args:
        model: 训练好的Transformer翻译模型
        src: 源序列张量
        src_mask: 源序列掩码
        max_len: 生成序列的最大长度
        beam_size: 束宽，建议取值2~5
        alpha: 长度惩罚系数，取值0~1，数值越大越倾向生成长句
        device: 运行设备
    Returns:
        得分最高的目标序列张量
    """
    model.eval()
    # 编码器输出全局复用
    memory = model.encode(src, src_mask)
    # 初始化候选束：每个元素为(累计得分, 当前序列, 已生成的序列长度)
    candidates = [(0.0, torch.tensor([BOS_IDX], dtype=torch.long, device=device), 1)]
    
    for _ in range(max_len - 1):
        next_candidates = []
        # 遍历所有当前候选
        for score, tgt, seq_len in candidates:
            # 如果当前候选已经结束，直接加入下一轮
            if tgt[-1].item() == EOS_IDX:
                next_candidates.append((score, tgt, seq_len))
                continue
            # 生成目标序列掩码
            tgt_mask = model.generate_square_subsequent_mask(seq_len).to(device)
            # 解码器推理
            out = model.decode(tgt.unsqueeze(1), memory, tgt_mask, src_mask)
            # 取最后一个token的概率分布
            prob = model.generator(out[-1])
            log_prob = torch.log_softmax(prob, dim=-1).squeeze(0)
            # 取topk个候选token
            topk_log_probs, topk_indices = torch.topk(log_prob, beam_size)
            
            # 生成新的候选
            for i in range(beam_size):
                token = topk_indices[i].item()
                new_score = score + topk_log_probs[i].item()
                new_seq = torch.cat([tgt, topk_indices[i].unsqueeze(0)], dim=0)
                new_seq_len = seq_len + 1
                # 应用长度惩罚
                penalized_score = new_score / (new_seq_len ** alpha)
                next_candidates.append((penalized_score, new_seq, new_seq_len))
        
        # 保留得分最高的beam_size个候选
        next_candidates.sort(key=lambda x: x[0], reverse=True)
        candidates = next_candidates[:beam_size]
        
        # 所有候选都已结束的话提前终止
        if all(cand[1][-1].item() == EOS_IDX for cand in candidates):
            break
    
    # 返回得分最高的序列
    return candidates[0][1]

使用方法 #

• 首先确认你代码中的特殊字符索引（BOS_IDX/EOS_IDX/PAD_IDX）与上述代码中的定义一致，不符的话直接修改代码开头的常量取值即可
• 替换官方教程中原greedy_decode的调用为上述beam_search_decode函数
• 可按需调整beam_size和alpha参数：束宽越大解码精度越高但推理速度越慢，普通场景设置为3即可获得明显优于贪心解码的效果
• 长度惩罚系数alpha建议取值0.6~1.0，需要生成更短的句子可以调低数值，需要生成长句可以调高数值

内容的提问来源于stack exchange，提问作者Purin Charunrattanapiriya