做LSTM序列训练的时候最烦的一件事不是模型收敛不了而是DataLoader在最后一个batch上给你报一个RuntimeError说我传进去的tensor形状不一致。我最早做文本分类的时候就撞过这个墙那会儿数据里有长有短的评论Dataloader默认的collate_fn直接对我做了stack长度不一致怎么都叠不起来。后来切到时间序列预测手里是不同站点的水文径流数据每个样本对应的时间窗长度也不一样同样的问题又冒出来了。可以说只要你的数据不是规规矩矩等长的LSTM训练的这条路上长度不一致就是绕不过去的坎。这篇文章就把我在这类问题上踩过的坑、用过的方案以及完整的代码路径理一遍。内容包括三种主流的处理方式重点讲清楚padding pack_padded_sequence这条组合拳怎么落地也会把DataLoader里collate_fn的写法、batch内的排序问题、模型的输出对齐这些容易被忽略的细节拆开讲直接给你一套能跑通的代码。适合正在用PyTorch做时间序列预测、自然语言处理或者任何以LSTM为baseline的变长序列任务的读者参考。1. 变长序列问题到底出在哪1.1 LSTM训练中的典型数据形态先说清楚“数据长度不一致”在实操里长什么样。最常见的场景有两类。第一类是NLP里的句子样本。一个batch里可能有“今天天气不错”这种短句也可能有几千字的文章分词或者转token之后长度差异很大。设计DataLoader的时候如果不特殊处理要求每个样本的形状完全一致本质上是要求一个batch内的tensor能够拼成一个矩形张量。第二类是时间序列预测。比如用LSTM做水文径流预报每个站点的观测时间窗可能不一样有的站点积累700个时刻有的只有340个做多个物理量联合输入时不同特征的时间覆盖范围也可能不一样。这种场景直接喂给DataLoader同样会挂在stack那一步。不光是原始的输入序列长度不齐标签也可能是对齐困难的。比如LSTM做序列到序列预测输入和输出各有自己的长度你要分别对它们做padding然后在计算loss的时候把padding位置遮蔽掉。这些操作看起来独立实际上一旦batch里的序列长度参差不齐所有环节都会被“长度”这一变量牵扯住。1.2 DataLoader在变长数据上的报错机制很多新手会问DataLoader不是自动帮我们整理batch吗为什么传list进去还会报错原因是DataLoader默认的collate_fn做得太“死”了。它内部会把一个batch里的样本按位置对齐然后交给torch.stack去拼接。torch.stack有一个硬性要求所有输入tensor的形状必须完全相同。比如你的样本形状是(34, 8)和(72, 8)stack没办法把它们拼成一个更高维的tensor因为对应维度对不上。我见过好几个版本的报错长这样RuntimeError: stack expects each tensor to be equal size, but got [34, 8] at entry 0 and [72, 8] at entry 1这个问题并不在模型而在数据装载阶段。解决方式也很直接你得自己写一个collate_fn接管从“一列长度不等的样本”到“一个batch”的转换逻辑。这篇博文所有的核心操作本质上都围绕这个自定义collate_fn来展开。2. 核心思路三条主流路线如何选处理变长序列业内常见的路线有padding pack、纯mask、长度分组动态padding。先别急着写代码这三条路线各有适合的场景选错了后面会很难受。2.1 路线Apadding pack_padded_sequence这是LSTM训练里最经典的做法。核心思想是先把一个batch内的所有序列pad到相同长度然后利用PyTorch自带的pack_padded_sequence把padding部分剔除掉让LSTM只对真实序列计算不浪费算力在无效的0上。这个方案的优点非常突出可以保留batch训练吞吐量高不改变LSTM的数学结构只要batch内长度排序正确结果就和逐条训练一致配合pad_packed_sequence后面取隐状态、算loss都能很自然地完成。我实际用下来这是对付长度不一致最稳定的方案。只要数据量不是极端不均衡比如一个batch里最长序列是1000最短只有1这个方案都能正常工作。2.2 路线Bmask机制处理mask的思路是在padding之后仍然让LSTM把所有padded位置都跑一遍但在计算loss或者最终输出时用mask把这些位置遮掉。LSTM本身不原生支持attention mask所以通常要靠手动构造一个(B, T)的bool矩阵在取输出或者算loss时做masked_select。这个方案适合实现相对简单的场景但有两个我比较介意的问题Transformer类模型可以依靠mask跳过无效位置的计算而LSTM做不到。哪怕padding位置的值是0LSTM还是会更新隐状态长序列里的padding会白白增加计算量而且可能引入无意义的隐状态污染。如果只是padding不pack隐状态在padding段依然会继续传递最后一个时间步如果恰好是padding你取到的输出很可能是废的。所以mask方案我更推荐在Transformer、或者用CNN做序列任务时使用LSTM场景下它算一个“能用但不够优雅”的降级方案。2.3 路线C长度分组动态batching第三条路线是分桶bucket策略按序列长度把样本分到不同组每个组内部选择相近长度的样本再进行padding。这样做可以让padding的浪费比例降到最低。这条路线特别适合序列长度分布悬殊、又想尽可能加速训练的场景。比如有的句子长200有的只有20如果硬把所有样本pool在一起padding一个batch可能90%都是padding。分桶之后长度为20左右的样本在一起长度为200左右的在一起batch内padding比例能大幅下降。代价是你需要自己写或者改造batch samplerDataLoader默认不会帮你做长度分组。torchtext的BucketIterator是现成实现用原生PyTorch的话我一般会自己写一个简单的batch_sampler按长度排序之后切块。具体代码在第4章给一个可用的版本。2.4 三条路线的对比总结我平时选型的逻辑很简单可以直接照搬方案优点缺点适合场景padding pack_padded_sequence计算高效结果精确配套成熟需要排序约束代码略繁绝大多数LSTM/RNN任务稳定首选mask机制实现直观概念简单padding段仍参与隐状态更新浪费算力长pad易污染状态Transformer、注意力模型或简单demo长度分组动态paddingpadding占比低训练快需要额外实现samplerbatch顺序不稳定长度分布极度不均衡的中大规模时间序列/NLP任务如果你刚起步直接走第一条路如果数据长度差异大到受不了再把第三条路加上去其实也就是在第一条路前面多了一个分桶环节。3. 实操用自定义collate_fn pack_padded_sequence解决这里是全文核心我直接给出一套完整可运行的方案。代码风格我尽量保持和实际项目一致不搞玩具demo你可以直接抄进自己的训练脚本里。3.1 自定义collate_fn的第一步paddingDataLoader在构建batch时会调用collate_fn把样本list合并成batch。默认的collate_fn是torch.utils.data._utils.collate.default_collate它会尝试对list内tensor做stack。我们要做的就是绕过它自己定义一个collateimport torch import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence def pad_collate(batch): # batch: list of dict每个dict包含 seq 和 target seqs [item[seq] for item in batch] targets torch.tensor([item[target] for item in batch]) # padding到batch内最大长度batch_firstTrue意味着输出形状为(B, T, D) seqs_padded pad_sequence(seqs, batch_firstTrue, padding_value0.0) lengths torch.tensor([len(seq) for seq in seqs], dtypetorch.long) return seqs_padded, lengths, targets这里关键函数是pad_sequence。它的作用是把一串长度不等的tensor在尾部补0直到对齐最长的那个。如果你不传batch_firstTrue默认输出维度是(T, B, D)容易绕晕我建议统一开batch_first。padding_value这个参数容易被忽视。默认是0对于大部分特征来说没问题。但如果你处理的是数值范围比较敏感的特征padding值会影响归一化统计建议padding之后再用mask把padding位置的数据从loss里剔除不要让它污染模型。3.2 第二步构造lengths并保持降序pack_padded_sequence有一个很重要的坑它默认要求batch内序列按长度降序排列也就是第一个样本最长最后一个最短。这个要求来自PyTorch内部对packed序列的布局设计。所以你在collate_fn里需要顺手做一个长度排序def pad_collate(batch): seqs [item[seq] for item in batch] targets torch.tensor([item[target] for item in batch]) lengths torch.tensor([len(seq) for seq in seqs], dtypetorch.long) sorted_indices torch.argsort(lengths, descendingTrue) seqs [seqs[i] for i in sorted_indices] targets targets[sorted_indices] lengths lengths[sorted_indices] seqs_padded pad_sequence(seqs, batch_firstTrue, padding_value0.0) return seqs_padded, lengths, targets注意排序不仅作用于输入也一定要同步作用于target。我最早写的时候只排了seq没排target结果模型训练了半个epoch才发现loss一直在震荡最后定位到是标签和样本对不齐那个教训挺惨的。这里额外说一句如果你对训练样本的顺序有特殊要求比如希望batch内顺序不被打乱可以在collate之后记录原始索引。但对于绝大多数训练任务顺序不影响最终结果你不需要执着于保持原顺序。3.3 第三步pack与LSTM前向DataLoader返回seqs_padded, lengths, targets之后模型forward里要做的操作是先pack再丢进LSTM再用pad_packed_sequence恢复。class LSTMClassifier(nn.Module): def __init__(self, input_size, hidden_size, num_layers, num_classes): super().__init__() self.lstm nn.LSTM( input_size, hidden_size, num_layersnum_layers, batch_firstTrue, bidirectionalFalse ) self.fc nn.Linear(hidden_size, num_classes) def forward(self, x, lengths): packed_x pack_padded_sequence( x, lengths.cpu(), batch_firstTrue, enforce_sortedTrue ) packed_out, (hn, cn) self.lstm(packed_x) out, out_lengths pad_packed_sequence(packed_out, batch_firstTrue) # 取每个样本最后一个有效time step的输出 idx (out_lengths - 1).unsqueeze(1).expand(-1, out.size(-1)).unsqueeze(1) last_output out.gather(1, idx).squeeze(1) logits self.fc(last_output) return logits几点要解释清楚。第一pack_padded_sequence接收的lengths必须是一个一维张量而且建议放在CPU上。如果直接传GPU张量某些版本的PyTorch会报错或者运行缓慢。我在新版PyTorch上也遇到过它内部调用需要CPU张量的情况所以安全起见统一.cpu()。第二enforce_sortedTrue意味着你必须确保长度降序。如果你不能保证排序可以设成enforce_sortedFalsePyTorch内部会自己重新排序。但这会引入额外开销且每个batch重新排序会打乱样本顺序我个人还是建议在collate_fn里手动排好序统一用True。第三hn的形状是(num_layers * num_directions, B, hidden_size)。如果你用的是双向LSTM最后一个输出的拼接方式要更小心。上面代码用的是单向LSTM直接取最后有效的输出。3.4 第四步从packed output中取有效信息LSTM的输出经过pad_packed_sequence之后padding位置的输出是0。这时候如果你直接取out[:, -1, :]取到的是最后一个time step而它大概率是padding位置输出没有任何意义。正确做法是根据out_lengths来取每个样本的最后一个有效time step。我在上面的代码里用了gather实际上还有更简单的写法last_output torch.stack([ out[i, out_lengths[i] - 1, :] for i in range(out.size(0)) ])如果样本量不大这个写法性能也能接受。gather版本更快但代码可读性差一点。训练时取hn的最后一层也可以这个后面在第4章专门讲一下区别。如果是序列到序列任务比如水文径流预报里用一个LSTM编码器输出每个时间步的预测值那么你在算loss时还要生成一个mask把padding位置的loss遮掉def sequence_mask(lengths, max_len): # lengths: [B], 返回 (B, max_len) bool矩阵True为有效位置 batch_size lengths.size(0) arange torch.arange(max_len, devicelengths.device) mask arange.expand(batch_size, max_len) lengths.unsqueeze(1) return mask mask sequence_mask(out_lengths, out.size(1)) # (B, T) mask mask.unsqueeze(-1) # 扩展到特征维度 loss F.mse_loss(out * mask, target * mask, reductionsum) / mask.sum()注意这里有另一个容易出错的地方out_lengths毕竟是模型自己算出来的在没排序的情况下会与输入顺序对不上。如果走的是enforce_sortedFalse最好直接使用pack_padded_sequence返回的重新排序后的索引或者干脆回到手动排序方案省心很多。4. 常见问题与排查技巧实录4.1 报错“ValueError: Expected batch_first...”这个报错几乎每个人都会遇到。原因是pack_padded_sequence和pad_packed_sequence都要关注batch_first参数而LSTM层也要设置batch_firstTrue。这三个地方必须保持一致否则就会在维度对不上时爆出奇怪的错误。我建议从一开始就统一走batch_firstTrue路线。这样输入形状就是(B, T, D)pack后是PackedSequencepad后恢复为(B, T, D)。如果你身边有老代码习惯用(T, B, D)也不是不行但混合两种风格非常容易出bug尤其在做时间序列预测时特征维度D、时间步T、batch维度B本来就容易绕再叠加上batch_first不统一基本就是噩梦。4.2 pack之后为什么输出全乱如果你在collate_fn里排好序了但是模型输出和标签对不上概率最大的原因是排序时只排了输入没排标签也忘了用enforce_sortedFalse。另一种情况是你在里面使用了enforce_sortedFalse然后取结果时用了原始输入的索引顺序去对齐。我见过一个比较隐蔽的坑PackedSequence本身携带了sorted_indices和unsorted_indices属性。如果你用了enforce_sortedFalse可以这样恢复原来的batch顺序packed_output, (hn, cn) self.lstm(packed_input) out, _ pad_packed_sequence(packed_output, batch_firstTrue) # 如果之前没手动排序需要用 unsorted_indices 恢复 original_indices packed_input.unsorted_indices out out[original_indices]这个属性建议记住排查乱序问题时很有用。不过我更推荐直接手动排序代码更可控。4.3 取hn还是取out.gather很多LSTM分类任务里取hn的最后一层应该就够了last_hidden hn[-1] # 单向LSTM取最后一层 logits self.fc(last_hidden)这种方式代码更简洁而且在大部分情况下结果和取最后有效时间步输出一致。因为序列最后一个真实time step的隐状态本来就会持续到LSTM结束。但有一个细节如果是多层LSTMhn[-1]才是最后一层hn[0]是第一层不要搞混如果是双向LSTM你需要自己拼接前向和后向的隐状态last_hidden torch.cat((hn[-2], hn[-1]), dim-1)这里的顺序不是绝对的要参考PyTorch文档对hn的定义。实际项目中我多是用out.gather方式因为它对双向LSTM、多层结构的兼容性更好逻辑也更直观。4.4 CPU/GPU上性能对比与加速建议很多人会问既然pad_sequence已经把序列padding了pack一下不是多此一举吗我实际在GPU上做过对比。一个batch里5条序列长度分别是300、50、40、30、20。如果不packLSTM会按最大长度300跑完所有序列等于是5条都跑了300步。pack之后LSTM只在5条序列的真实范围内递归省掉的计算量非常可观。在GPU上pack的收益不只是计算量还在于显存占用。不pack时每个time step的中间输出都要保存用于反向传播尺寸是(B, T, hidden)pack后是总有效长度显存能省不少。不过要注意pack之后输入数据不再是连续张量有些GPU算子可能没针对PackedSequence做优化极端情况下反而比直接padding慢。我的经验是当batch内最大长度和平均长度差距超过2倍时pack的收益非常明显差距很小的时候比如大多数样本长度都在25到30之间那直接padding不用pack性能差别也不大代码还能少写几行。4.5 常见问题速查我把自己多次遇到的坑整理成一个速查表排查时直接对照表现可能原因处理方式RuntimeError: stack expects each tensor to be equal size没有自定义collate_fn默认在stack用pad_sequence代替stackpack_padded_sequence报length order错误batch内序列未按长度降序排列在collate_fn里argsort(descendingTrue)或设置enforce_sortedFalse输出全乱accuracy震荡排序后target没有跟着输入一起排排序时同步对target、索引做相同操作取最后时刻输出全是0直接用out[:, -1, :]取到了padding位置用out_lengths - 1索引去gatherloss在padding位置乱跳算loss时未屏蔽padding位置用sequence_mask生成maskloss只在有效位置计算显存爆了序列过长且padding比例高分桶batch_sampler或缩短max_len训练结果依赖batch内顺序使用了enforce_sortedFalse且输出未恢复顺序恢复unsorted_indices或手动排序4.6 分桶batch_sampler的一个可用实现如果长度分布悬殊我给一个最简单可用的分桶实现配合第3章的pad_collate一起用。from torch.utils.data import Sampler class BucketBatchSampler(Sampler): def __init__(self, lengths, batch_size, shuffleTrue): self.lengths lengths self.batch_size batch_size self.shuffle shuffle self.batches self._make_batches() def _make_batches(self): indices list(range(len(self.lengths))) if self.shuffle: # 打乱后按长度排序得到长度相近的样本索引块 import random random.shuffle(indices) indices.sort(keylambda i: self.lengths[i]) return [indices[i:i self.batch_size] for i in range(0, len(indices), self.batch_size)] def __iter__(self): if self.shuffle: random.shuffle(self.batches) return iter(self.batches) def __len__(self): return len(self.batches)使用方式和普通batch_sampler一致传入DataLoader的batch_sampler参数同时去掉batch_size参数。要注意的是使用batch_sampler后DataLoader不会再调用默认的batch采样所以collate_fn里不需要再对batch选择负责只需要处理padding和排序。这个实现还有个坑每个batch内部仍然是原始顺序所以需要pad_collate里继续做长度降序。因为分桶只是相近长度放在一起不代表一定严格降序。5. 经验心得与扩展建议5.1 关键操作顺序把完整流程再串一遍方便你照着搭在Dataset的__getitem__里每条样本返回一个原始长度的tensor不padding不截断。在collate_fn里先取每个样本长度按降序排列。用pad_sequence把排序后的序列padding到batch内最大长度。把lengths转成CPU张量随batch一起返回。模型forward里先用pack_padded_sequencepack再进入LSTM。输出用pad_packed_sequence恢复形状。根据out_lengths取有效最后一步或者用hn[-1]。计算loss时用mask屏蔽padding位置。这套顺序我已经在好几个项目里复用基本稳定。核心原则就是能到集装线tensor操作之前处理好的变长问题绝不要拖到模型内部再纠结。5.2 处理极端长度差异的隐藏技巧如果你手里的序列长度极端比如一个batch里同时出现10个time step和5000个time step的数据即使pack了短序列的梯度也会很弱长序列可能要跑很久。我习惯再叠加两层保险。保险一是设置最大长度上限。时间序列任务里可以按训练数据的长度统计切分比如取99分位作为截断阈值超过的直接截断不要保留几个极端长样本拖累整个训练速度。文本任务里truncation到512或者128都是常规操作。保险二是长度分桶。第4章给的BucketBatchSampler虽然简单但效果立竿见影。分桶后每个batch内长度方差变小padding比例下降训练速度能提升好几倍。我在一个径流预报任务里原本一个epoch要跑12分钟分桶后只需要5分钟模型效果还因为padding干扰减少略有提升。5.3 为什么我不用PackedSequence做数据增强有一段时间我想过既然pack能跳过padding那是不是可以在这个基础上做序列数据的随机裁剪、拼接增强实际操作后发现PackedSequence对序列长度变化很敏感一旦你在batch内部做了裁剪每个样本的原始长度变化了packed结构也要重建比较繁琐。所以我的建议是数据增强尽量在Dataset层面做也就是在__getitem__返回之前就完成裁剪、加噪、归一化等操作。collate_fn和模型前向只负责“整理长度”不要在里面塞复杂的业务逻辑否则后期维护会很难受。5.4 关于序列排序的一点个人观点很多教程都会说“pack前必须按长度降序排列”但实际使用中我见过有人直接设置enforce_sortedFalse省掉了排序逻辑。这个做法没问题但我要提醒一句如果你用了enforce_sortedFalsePyTorch会在每个batch内部重新排序意味着排序逻辑在每次前向都执行一次多少有额外开销。更关键的是如果模型里有dropout、batch normalization这些随机性组件重新排序不会影响训练结果但如果你的batch间顺序敏感比如做课程学习curriculum learning还是自己控制排序更稳。我平时默认就是手动排序不依赖enforce_sortedFalse长期跑下来没有出现过问题。排序代码也不难写就三五行风险可控。5.5 后续还可以怎么优化只要把数据加载这块理顺了后面要提速还可以考虑几个方向。一是用torch.compile加速LSTM模型但要注意PackedSequence和torch.compile的兼容性我目前试下来的经验是新版本PyTorch能跑但并不是所有算子都达到加速预期。二是把pack操作放进collate_fn之外改为在dataset里预先计算长度和排序索引减少重复排序。三是对超长序列做segment划分把长序列切成多个短段每个段分别过模型再接state这是另一个话题了但思想上和变长处理也是一脉相承的。我个人在实际项目里的体会是变长序列处理这个事其实不复杂难的是所有细节都对齐。那些看起来不重要的排序、lengths放CPU、batch内顺序同步如果你忽略任何一个训练过程都会用诡异的报错或结果反哺你。把上面这套流程跑通之后你会发现自己以后遇到任何需要处理不定长序列的任务都可以先在DataLoader这一层把长度问题解决掉而不是在模型后面打补丁。这就是这篇文章最想传达的经验。