pytorch LSTM处理变长序列代码

序列数据一般长度是不固定的，所以在训练LSTM这类模型时，一个batch内的数据需要padding到相同的长度，并在训练中忽略相关影响。

整理了pytorch LSTM利用pack_padded_sequence、pad_packed_sequence处理变长序列数据代码：

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

# 假设我们有一个批次大小为3的序列数据，序列长度分别为[3, 2, 4], 然后用0填充使它们长度一致
data = [
    [1, 2, 1, 0],
    [1, 2, 0, 0],
    [1, 2, 3, 4]
]
# 将数据转换为PyTorch张量
data_tensor = torch.tensor(data)

# 真实的序列长度
lengths = [3, 2, 4]

# 首先按照序列长度降序排列数据
lengths, perm_index = torch.sort(torch.tensor(lengths), descending=True)
data_tensor = data_tensor[perm_index]

# 定义一个Embedding层，表示序列每个元素
emb = nn.Embedding(num_embeddings=5, embedding_dim=2)
# 定义一个LSTM
lstm = nn.LSTM(input_size=2, hidden_size=2, batch_first=True)

# 使用pack_padded_sequence打包序列，需要提供数据和实际长度
packed_input = pack_padded_sequence(emb(data_tensor), lengths, batch_first=True)
# 传递打包的序列到LSTM
packed_output, (hidden, cell) = lstm(packed_input)

# 使用pad_packed_sequence做一个padding操作还原到输入形状
total_length = data_tensor.size(1)
output, _ = pad_packed_sequence(packed_output, batch_first=True, total_length=total_length)