PyTorch中LSTM网络填充与打包序列时的问题

问题 #

我正在构建一个简单的LSTM神经网络，用PyTorch的Dataset和DataLoader处理时序数据，拆分序列与批次。为解决最后一批数据长度不一致的问题，我用了填充（padding）和打包（packing）操作。

我的collate_fn函数 #

def collate_data(batch):
    sequences, targets = zip(*batch)
    
    lens = [len(seq) for seq in sequences]
    print(f"Lens before padding: {lens}")

    padded_seq = pad_sequence(sequences=sequences,batch_first=True,
    padding_value=float(9.99e10))

    print(f"Lens after padding: {[len(seq) for seq in padded_seq]}")

    padded_targets = pad_sequence(sequences=targets,batch_first=True,
    padding_value=float(9.99e10))

    packed_batch=pack_padded_sequence(padded_seq,lengths=lens,batch_first=True,\
    enforce_sorted=False)

    print(f"Packed batch lengths: {packed_batch.batch_sizes}")

    return packed_batch, padded_targets

我的forward方法 #

def forward(self,x ):
    lstm = self.lstm
    batch_size = self.batch_size

    h0 = torch.zeros(self.num_layers,batch_size,self.hidden_size,)   
    c0 = torch.zeros(self.num_layers,batch_size,self.hidden_size,)

    packed_lstm_out, (hn,cn) = lstm(x, (h0,c0))
    
    print(f"lstm_out size: {packed_lstm_out.data.size}")        
    unpacked_lstm_out = unpack_sequence(packed_sequences=packed_lstm_out,)        
    print(f"Unpacked lengths: {[len(seq) for seq in unpacked_lstm_out]}")

    unpacked_lstm_tensor = torch.stack(unpacked_lstm_out,dim=0).float().\
    requires_grad_(True)

    print(unpacked_lstm_tensor.shape)

    output = self.fc1(unpacked_lstm_tensor[:,-1,:])

    return output

问题现象 #

执行torch.stack(unpacked_lstm_out, dim=0)时报错，提示张量尺寸不一致，且仅在最后一批（已填充的批次）出现该问题。最后一批的打印输出如下：

Lens before padding: [10, 10, 10, 10, 10, 10, 10, 9, 8, 7, 6, 5]
Lens after padding: [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
Packed batch lengths: tensor([12, 12, 12, 12, 12, 11, 10,  9,  8,  7])
lstm_out size: torch.Size([105, 16])
Unpacked lengths: [10, 10, 10, 10, 10, 10, 10, 9, 8, 7, 6, 5]

我怀疑问题出在pack_padded_sequence()的使用上，但不知道具体原因和解决方法，请问如何修复让forward函数解包后的张量尺寸一致？

解决方案 #

问题根源在于unpack_sequence返回的是长度不一的张量列表，直接用torch.stack会因为尺寸不匹配报错。以下是两种可行的修复方式：

方法一：用pad_packed_sequence直接解包（推荐） #

不要用unpack_sequence，改用pad_packed_sequence直接得到填充对齐后的张量，同时会返回每个序列的真实长度：

def forward(self, x):
    lstm = self.lstm
    # 动态获取当前批次的batch_size，避免固定值导致最后一批不匹配
    batch_size = x.batch_sizes[0].item()
    num_layers = self.num_layers
    hidden_size = self.hidden_size

    # 确保h0、c0和输入数据在同一设备（CPU/GPU）
    h0 = torch.zeros(num_layers, batch_size, hidden_size, device=x.data.device)   
    c0 = torch.zeros(num_layers, batch_size, hidden_size, device=x.data.device)

    packed_lstm_out, (hn, cn) = lstm(x, (h0, c0))
    
    # 直接解包为填充后的张量，batch_first=True对应输入格式
    unpacked_lstm_tensor, lengths = torch.nn.utils.rnn.pad_packed_sequence(packed_lstm_out, batch_first=True)
    
    # 取每个序列的最后有效步输出（根据真实长度索引）
    last_indices = lengths - 1
    output = self.fc1(unpacked_lstm_tensor[torch.arange(batch_size), last_indices, :])

    # 也可以直接用hn[-1]，效果一致——LSTM的hn就是每个序列最后一步的隐藏状态
    # output = self.fc1(hn[-1])

    return output

方法二：解包后手动填充对齐 #

如果一定要用unpack_sequence，可以在解包后对每个序列重新填充，统一到当前批次的最大长度：

def forward(self, x):
    lstm = self.lstm
    batch_size = x.batch_sizes[0].item()
    num_layers = self.num_layers
    hidden_size = self.hidden_size

    h0 = torch.zeros(num_layers, batch_size, hidden_size, device=x.data.device)   
    c0 = torch.zeros(num_layers, batch_size, hidden_size, device=x.data.device)

    packed_lstm_out, (hn, cn) = lstm(x, (h0, c0))
    
    # 解包得到长度不一的序列列表
    unpacked_lstm_out = unpack_sequence(packed_lstm_out)
    # 获取当前批次的最大序列长度
    max_len = max(len(seq) for seq in unpacked_lstm_out)
    # 对每个序列进行填充，对齐到max_len
    padded_unpacked = [torch.nn.functional.pad(seq, (0,0,0, max_len - len(seq))) for seq in unpacked_lstm_out]
    # 现在可以正常stack
    unpacked_lstm_tensor = torch.stack(padded_unpacked, dim=0).float()

    # 取每个序列的最后有效步输出（需结合真实长度，不能直接取-1，避免取到填充值）
    # 这里需要从collate_fn传递真实长度，或者用x.batch_sizes计算，更简单的方式是直接用hn[-1]
    output = self.fc1(hn[-1])

    return output

额外注意事项 #

1. 不要用固定的self.batch_size：最后一批的batch_size可能小于预设值，必须从输入动态获取（比如x.batch_sizes[0].item()），否则h0和c0的尺寸会不匹配。
2. 设备一致性：h0和c0要和输入数据在同一个设备上（CPU/GPU），否则会出现张量设备不匹配的错误。
3. 填充值不影响计算：你用的9.99e10填充值不会被LSTM计算，因为pack_padded_sequence会自动跳过padding部分，这部分无需担心。

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