变长输入的Mini Batch训练实现问题（PyTorch RNN场景）

嘿，我完全懂你现在的卡点——刚摸PyTorch的RNN时，从单样本SGD转到小批量训练，尤其是变长序列的情况，确实容易懵。咱们一步步拆解，把逻辑理清楚，你很快就能搞定

你现在手里的两个列表：一个是变长LongTensor序列，一个是对应标签。要做小批量，首先得把这些数据包装成PyTorch的Dataset，再用DataLoader加载。但因为序列长度不一样，得先做padding补全，同时记录每个序列的原始长度（后面模型要用到）。

先写个自定义的Dataset：

import torch
from torch.utils.data import Dataset, DataLoader

class SequenceDataset(Dataset):
    def __init__(self, sequences, labels):
        self.sequences = sequences
        self.labels = labels

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        # 返回单条序列和对应标签
        return self.sequences[idx], self.labels[idx]

然后关键是写collate_fn——因为默认的DataLoader处理不了变长序列，这个函数用来给每个batch的序列补全到统一长度，同时记录原始长度：

def collate_fn(batch):
    # batch是列表，每个元素是(sequence, label)
    sequences, labels = zip(*batch)
    
    # 1. 记录每个序列的原始长度
    seq_lengths = torch.tensor([len(seq) for seq in sequences], dtype=torch.int64)
    
    # 2. 找到当前batch里最长的序列，作为补全后的长度
    max_seq_len = max(seq_lengths)
    
    # 3. 给每个序列补0（你也可以用其他padding值，比如词汇表外的索引）
    padded_seqs = torch.zeros(len(sequences), max_seq_len, dtype=torch.long)
    for i, seq in enumerate(sequences):
        padded_seqs[i, :len(seq)] = seq
    
    # 4. 把标签转成Tensor
    labels = torch.tensor(labels, dtype=torch.long)
    
    # 返回补全后的序列、原始长度、标签
    return padded_seqs, seq_lengths, labels

最后创建DataLoader：

# 假设你的sequences是LongTensor列表，labels是标签列表
dataset = SequenceDataset(sequences, labels)
# 这里batch_size设成你想要的大小，比如32
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)

之前单样本直接喂序列就行，但批量有padding的话，模型会把padding部分也算进去，影响最终隐藏状态。这时候得用pack_padded_sequence和pad_packed_sequence来跳过padding的计算。

举个GRU的例子（LSTM逻辑类似，只是隐藏状态是元组）：

import torch.nn as nn

class SeqClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.gru = nn.GRU(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, num_classes)

    def forward(self, padded_seqs, seq_lengths):
        # 1. 先过embedding层
        embeds = self.embedding(padded_seqs)  # shape: [batch_size, max_seq_len, embed_dim]
        
        # 2. 必须先按序列长度降序排序！pack_padded_sequence要求这么做
        sorted_lengths, sorted_idx = seq_lengths.sort(descending=True)
        sorted_embeds = embeds[sorted_idx]
        
        # 把补全后的序列打包，跳过padding部分的计算
        packed_embeds = nn.utils.rnn.pack_padded_sequence(sorted_embeds, sorted_lengths, batch_first=True)
        
        # 3. 喂给GRU
        packed_output, hidden = self.gru(packed_embeds)
        # GRU的hidden是最后一步的隐藏状态，取最后一层即可：shape [batch_size, hidden_dim]
        final_hidden = hidden[-1, :, :]
        
        # 4. 把隐藏状态恢复回原batch的顺序，不然标签和预测会对应不上
        _, unsorted_idx = sorted_idx.sort()
        final_hidden = final_hidden[unsorted_idx]
        
        # 5. 全连接层做分类
        logits = self.fc(final_hidden)
        return logits

这里划重点：

• 排序是必须的，pack_padded_sequence要求序列按长度从长到短排列，不然会报错
• 排序后要把结果恢复原顺序，不然每个样本的预测和标签会错位
• 如果是LSTM，hidden是(h_n, c_n)，取h_n[-1]作为最终隐藏状态就行

之前单样本的循环改成从DataLoader里取batch就行：

import torch.optim as optim

# 初始化模型、损失、优化器（替换成你的实际参数）
vocab_size = 1000  # 你的词汇表大小
embed_dim = 128
hidden_dim = 256
num_classes = 5  # 你的分类类别数

model = SeqClassifier(vocab_size, embed_dim, hidden_dim, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 训练循环
model.train()
for epoch in range(10):  # 替换成你的epoch数
    total_loss = 0.0
    for padded_seqs, seq_lengths, labels in dataloader:
        optimizer.zero_grad()
        
        # 前向传播
        logits = model(padded_seqs, seq_lengths)
        
        # 计算损失
        loss = criterion(logits, labels)
        
        # 反向传播+更新参数
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    print(f"Epoch {epoch+1}, 平均损失: {total_loss/len(dataloader):.4f}")

• padding值的选择：如果你的序列是词汇索引，padding值要选词汇表外的（比如0，前提是你没把0用作有效词索引），不然embedding层会把padding当成有效词处理
• batch_first参数：RNN的batch_first=True要和pack_padded_sequence的batch_first=True保持一致
• 隐藏状态维度：GRU/LSTM的hidden维度是[num_layers*num_directions, batch_size, hidden_dim]，取最后一层[-1, :, :]就是每个样本的最终隐藏状态

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