PyTorch中Seq2Seq模型的批量处理机制及实现疑问

我之前也踩过这个坑！当你处理变长序列的批量数据时，光做填充还不够——模型会把填充的无效部分当成有效输入计算，这肯定会导致输出异常。下面是几个关键的修复点：

1. 给编码器传入真实序列长度 #

不管你用RNN、LSTM还是GRU，PyTorch的这类层都支持通过pack_padded_sequence处理变长序列，让模型只计算真实有效部分：

import torch
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# 假设编码器是LSTM
encoder_lstm = torch.nn.LSTM(input_size=encoding_dimension, hidden_size=hidden_size, batch_first=True)

# 先对序列长度降序排序（pack_padded_sequence强制要求）
sorted_lengths, sorted_idx = torch.sort(torch.tensor(sequence_lengths), descending=True)
sorted_batch = batch[sorted_idx]

# 打包序列，跳过填充部分
packed_input = pack_padded_sequence(sorted_batch, sorted_lengths, batch_first=True)
packed_output, (hidden, cell) = encoder_lstm(packed_input)

# 可选：解包回填充后的形状（如果需要后续处理）
output, _ = pad_packed_sequence(packed_output, batch_first=True)

这里要注意，处理完后如果需要恢复原batch顺序，记得用sorted_idx反向索引还原。

2. 给注意力层添加编码器掩码（如果用了注意力） #

如果你的Seq2Seq带注意力机制（比如Bahdanau或Luong注意力），必须生成掩码让注意力层忽略填充位置：

# 生成编码器掩码：形状[batch_size, max_seq_len]，有效位置为True，填充位置为False
encoder_mask = torch.zeros(batch_size, max_seq_len, dtype=torch.bool)
for i in range(batch_size):
    encoder_mask[i, :sequence_lengths[i]] = True

# 以PyTorch原生MultiheadAttention为例，传入掩码
attention = torch.nn.MultiheadAttention(embed_dim=hidden_size, num_heads=4, batch_first=True)
attn_output, attn_weights = attention(
    query=decoder_output, 
    key=encoder_output, 
    value=encoder_output, 
    key_padding_mask=~encoder_mask  # 注意这里取反，因为该参数是"需要被mask的位置为True"
)

如果是自定义注意力，要在计算注意力分数时把填充位置设为负无穷，确保softmax后这些位置权重为0：

# 假设attn_scores形状是[batch_size, dec_seq_len, enc_seq_len]
attn_scores.masked_fill_(~encoder_mask.unsqueeze(1), -1e9)
attn_weights = torch.softmax(attn_scores, dim=-1)

3. 解码器的输入处理与掩码 #

解码器的目标序列同样是变长的，需要两步处理：

• 对目标序列填充后，同样用pack_padded_sequence传入RNN类解码器
• 生成下三角掩码防止解码器看到未来token，同时结合填充掩码忽略无效位置

# 生成防止未来信息泄露的下三角掩码
def create_target_mask(target_seq_len):
    return torch.tril(torch.ones(target_seq_len, target_seq_len)).bool()

# 结合填充掩码：target_mask_pad是目标序列的填充掩码（[batch_size, target_seq_len]）
target_mask = create_target_mask(target_seq_len).unsqueeze(0).repeat(batch_size, 1, 1)
target_mask = target_mask & target_mask_pad.unsqueeze(1)

4. 损失计算时忽略填充token #

最后计算损失时，一定要排除填充位置的损失，否则会干扰模型训练：

# 假设填充对应的token id是0
criterion = torch.nn.CrossEntropyLoss(ignore_index=0)
# 把输出和目标展平后计算损失
loss = criterion(output.view(-1, vocab_size), target.view(-1))

把这些步骤补上，应该就能解决你遇到的输出异常问题了！

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