第一次把nn.GRU塞进模型里的时候我并没有被什么门控公式难住反而是被它那套输入输出形状约定绊了一跤。当时我按(batch, seq, feature)的顺序把张量喂进去程序不报错loss 也照常往下掉但验证集指标始终比预期差一截排查了两天才发现是自己把seq_len和batch两个维度搞反了——偏偏那两个数字刚好都是 32广播机制把错误完全掩盖了。这篇就把torch.nn.GRU的输入和输出从头拆一遍input、h_0、output、h_n这四个张量分别是什么形状多层和双向配置会把形状搅成什么样变长序列为什么必须做 pack以及训练循环里那些文档不会明说、但踩一次就够的坑。文中代码都能直接复制运行适合已经会写 PyTorch 训练循环、但每次用到 GRU 都要回去翻文档的人。1. 从一次维度错位说起torch.nn.GRU 输入端的四条硬性约定1.1 input 的每个维度分别管什么构造 GRU 时只有两个参数是必填的input_size和hidden_size。hidden_size好理解就是隐状态的宽度input_size最容易被误读它描述的是单个时间步上、单个样本的特征维度既不是序列长度也不是 batch 大小。举个具体的例子你有一批句子先通过nn.Embedding转成向量得到形状(N, L, E)的张量此时喂给 GRU 的input_size应该是E而不是L也不是N。我在早期项目里就见过同事把input_size写成词表大小vocab_size训练能跑但收敛极慢原因就是参数矩阵的第一维被凭空放大了几万倍。默认顺序是(seq_len, batch, input_size)也就是时间维排在最前面。为什么不是直觉上的 batch 优先因为 PyTorch 的 RNN 系列最早是对着底层高性能实现对齐的那里时间维天然就在第 0 维框架为了少一次转置就沿用了这个约定。想要 batch 优先得显式打开batch_firstTrue。1.2 batch_first 只管两头不管中间这里有个反直觉的细节也是我见过最多的记混点batch_first只影响input和output完全不影响h_0和h_n。隐状态的 batch 维永远固定在第 1 位。张量batch_firstFalse默认batch_firstTrueinput(L, N, H_in)(N, L, H_in)output(L, N, D*H_out)(N, L, D*H_out)h_0(D*num_layers, N, H_out)同左不受影响h_n(D*num_layers, N, H_out)同左不受影响表里的L是序列长度N是 batchD是方向数单向为 1双向为 2H_out就是hidden_size。提示如果你手写h_0时用(N, H)两维传进去会直接报维度错误写成(num_layers, N, H)在单向情况下能跑但双向配置下又会对不上。最稳的写法是torch.zeros(num_layers * num_directions, N, hidden_size)把这几个量算清楚再传。1.3 h_0 可以不给但给了就得全对h_0是可选参数不传的话模块内部会自动填零。但只要传就必须同时满足三个条件形状对、dtype和模块权重一致、设备一致。第三个最容易出事——模型搬到 GPU 上h_0还留在 CPU报错信息是设备不匹配但错误栈会指向 forward 内部第一眼看不出是自己构造的隐状态的问题。还有一个隐蔽的坑是张量的内存连续性。如果你先做了x.transpose(0, 1)再取某一段当输入得到的张量在内存里不是连续的某些后端路径会直接拒绝抛出类似 Expected hidden to be contiguous 的提示。解决办法很便宜.contiguous()一下就行。2. output 与 h_n 的关系亲手跑一遍比背公式管用2.1 单向单层配置下两者数值完全重合nn.GRU的forward返回的是一个二元组(output, h_n)。很多人第一次看到两个张量直觉会以为一个是汇总一个是细节其实它们的区别只在时间维度的取值方式上output保留了每一个时间步的信息h_n只保留每个层、每个方向在最后一步的信息。跑一段最短的验证代码import torch import torch.nn as nn torch.manual_seed(0) gru nn.GRU(input_size10, hidden_size20) x torch.randn(5, 3, 10) # (L5, N3, H_in10) out, h gru(x) print(out.shape) # torch.Size([5, 3, 20]) print(h.shape) # torch.Size([1, 3, 20]) print(torch.allclose(out[-1], h[-1])) # Trueout[-1]是第 4 个时间步的输出h[-1]是第 0 层唯一一层处理完第 4 个时间步之后的隐状态。单向情况下它们物理上是同一份数据我一般用torch.allclose而不是来判断因为极少数情况下不同算子融合路径会带来浮点级的尾差。2.2 h_n 的堆叠顺序是先层后方向h_n的第 0 维长度等于num_layers * num_directions但索引规则不是所有正向排前面、所有反向排后面而是层优先第 0 层占前两个位置正向在前、反向在后第 1 层再占后面两个以此类推。写代码时可以用h_n.view(num_layers, num_directions, N, H)把它重排成一个更直观的四维张量我个人在做多层双向模型时几乎都会加这一步省得每次都要在脑子里数下标。2.3 output 永远是最后一层的输出这点要记牢一个常见的误解是以为output把所有层的输出都堆在一起了实际不是——output只包含最后一层在每个时间步上的输出。中间层的结果只体现在h_n里。所以当你需要中间层的表示比如做多层特征融合、给不同层加辅助损失时只能一层一层手动调用或者用nn.GRU之外的方式拆开。顺带一个很实用的等式在单向不管多少层的情况下h_n[-1]永远等于output在最后一个时间步上的切片。多层单向时output[-1]依然是最后一层的最后一步所以这个等式始终成立。这也是为什么很多分类代码里out[-1]和h[-1]可以互换使用。3. 多层与双向开关一开形状就开始连锁变化3.1 num_layers 只撑大 h_n不改 output 的宽度把num_layers从 1 改成 2output的最后一维仍然是hidden_size因为只有最后一层会输出但h_n的第 0 维从 1 变成 2。中间层的输入维度由上一层自动衔接第一层吃input_size之后的层都吃hidden_size这些都不用你操心。真正需要操心的是如果你想把h_n从第一层传递到下一批数据继续用得注意h_n[0]是第一层的状态、h_n[-1]是最后一层的状态方向别搞反。3.2 bidirectional 把输出宽度直接翻倍打开bidirectionalTrue之后D变成 2output的最后一维变成2 * hidden_size。拼接顺序是正向在前、反向在后前hidden_size维来自正向后hidden_size维来自反向。gru nn.GRU(10, 20, num_layers2, bidirectionalTrue) x torch.randn(5, 3, 10) out, h gru(x) print(out.shape) # torch.Size([5, 3, 40]) print(h.shape) # torch.Size([4, 3, 20]) 两层 x 两方向注意h_n的最后一维不会翻倍仍然是hidden_size翻倍只发生在output上。因为正向和反向各自维护一套独立的隐状态只是在输出时被拼到了一起。3.3 反向那一路的最后状态其实对应第一个时间步这是双向 GRU 里最值得单独拿出来讲的一点。反向这一路是从序列末尾往前处理的所以它处理完所有输入之后的隐状态落在output的第 0 个时间步的后半段而不是最后一个时间步。验证一下gru nn.GRU(10, 20, bidirectionalTrue) out, h gru(torch.randn(5, 3, 10)) print(torch.allclose(out[-1, :, :20], h[0])) # True正向的终态在末尾 print(torch.allclose(out[0, :, 20:], h[1])) # True反向的终态在开头由此可以推出一个很实用的结论如果你要从双向 GRU 里提取看过整句话之后再给出的句子表示最优取法是把output[-1]的前半段和output[0]的后半段拼起来因为这两段分别是正向、反向各自看完整个序列之后的表示。而output[-1]的后半段只见过最后一个词信息量很少。如果懒得做这套拼接最稳的做法是直接对output做池化sent_repr out.mean(dim0) # (N, 40) # 或者只用两侧终态拼接 sent_repr torch.cat([h[0], h[1]], dim-1) # (N, 40)后者在分类任务里用得最多因为它取的是两个方向真正的收敛态不用管时间维下标怎么数。4. 变长序列padding 之后不做 pack模型在悄悄学填充符4.1 不 pack 到底错在哪这是我认为 GRU 使用中危害最大、又最不容易被发现的一个问题。假设一个 batch 里有三条长度分别为 5、2、3 的序列为了凑成矩阵你必须把短的补到 5。如果你直接把补零之后的张量喂给 GRU那么对于第二条序列GRU 在第 5 个时间步之后的隐状态是它读过三个填充符之后的状态而不是读完第二个真实词之后的状态。此时h_n里装的已经不是句子的语义了。更糟的是这件事不会报错。模型照样训练指标可能只是略微变差你会以为是数据或者超参数的问题很难定位到这里。4.2 从 padding 到 pack 的完整链路标准做法是用pack_padded_sequence把填充部分折叠掉让 GRU 只在真实长度上计算输出再用pad_packed_sequence还原成带填充的矩阵。import torch import torch.nn as nn from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence emb nn.Embedding(100, 32, padding_idx0) gru nn.GRU(32, 64, batch_firstTrue) x torch.tensor([ [1, 2, 3, 4, 5], [6, 7, 0, 0, 0], [8, 9, 10, 0, 0], ]) lengths torch.tensor([5, 2, 3]) # 必须是 CPU 上的 int64 e emb(x) # (3, 5, 32) packed pack_padded_sequence(e, lengths, batch_firstTrue, enforce_sortedFalse) out_packed, h gru(packed) out, out_lengths pad_packed_sequence(out_packed, batch_firstTrue) print(out.shape) # torch.Size([3, 5, 64]) print(h.shape) # torch.Size([1, 3, 64]) print(out_lengths) # tensor([5, 2, 3])几个容易忽略的点第一PackedSequence内部永远是时间维优先的所以即使模块设了batch_firstTrue从 pack 得到的输出仍然要先pad_packed_sequence才能按 batch 维处理。第二喂进去的是PackedSequence那output出来也是PackedSequence不能直接.mean(dim0)必须先还原。第三padding_idx0在 embedding 里设上能进一步减少填充符带来的干扰虽然 pack 之后它本来就参与不了计算但万一某段逻辑绕过了 pack这行配置能兜底。4.3 enforce_sorted 与三类高频报错pack_padded_sequence默认要求序列长度是降序排列的如果不满足又不显式设置enforce_sortedFalse就会得到一个提示长度未排序的错误。传enforce_sortedFalse之后函数内部会自动排序、计算、再还原顺序代价是一次额外的索引操作绝大多数场景完全可以接受。日常最容易撞上的三类报错我整理成了表报错关键词根本原因处理方式lengths must be a CPU tensor长度张量被放到 GPU 上了构造时不要.to(device)保持在 CPUlengths must be of type torch.int64用了默认的 int32 或浮点构造时写dtypetorch.longsorted_indices / decreasing order长度没降序且未关排序检查加enforce_sortedFalseExpected all tensors on same device长度或 h_0 与权重设备不一致统一 device.contiguous()兜底还有一个替代方案不 pack改为手动按真实长度取output。具体是构造index (lengths - 1).view(-1, 1, 1).expand(-1, 1, hidden_size)再output.gather(0, index)。这种写法在自定义层数解耦、需要逐层干预的场合更灵活但速度上不如 pack我在序列不算长的项目里才会偶尔用它。5. 把 GRU 接到下游任务分类头和序列标注头的写法完全不同5.1 句子分类拿隐状态接一个线性层最典型的结构是 embedding 加 GRU 加线性分类头。核心是决定用哪一份张量作为句子表示。class GruClassifier(nn.Module): def __init__(self, vocab_size, emb_dim, hidden, num_cls): super().__init__() self.emb nn.Embedding(vocab_size, emb_dim, padding_idx0) self.gru nn.GRU(emb_dim, hidden, num_layers1, batch_firstTrue, bidirectionalTrue) self.fc nn.Linear(hidden * 2, num_cls) def forward(self, x, lengths): e self.emb(x) # (N, L, E) packed pack_padded_sequence(e, lengths, batch_firstTrue, enforce_sortedFalse) _, h self.gru(packed) # (2, N, H) feat torch.cat([h[0], h[1]], dim-1) # (N, 2H) return self.fc(feat)这里用h[0]和h[1]而不是output[-1]原因是双向配置下output[-1]只含正向的收敛态反向那半段几乎是没用的。这一点和单向模型的写法差别很大改配置时千万别只改bidirectional忘了改取法。5.2 序列标注必须用完整的 output做词性标注、实体识别这类逐时间步输出的任务时h_n是完全不够用的因为每一个时间步都要出结果。写法是还原后的output直接过一层Linearlogits self.fc(out) # (N, L, num_tags)这里有一个细节值得注意还原后的output在填充位置上的值是 GRU 对填充符计算出来的结果不是零。所以计算损失时必须配合 mask把填充位置排除掉mask (x ! 0) # (N, L) loss criterion(logits.transpose(1, 2), y) loss (loss * mask).sum() / mask.sum()如果直接在整段序列上算平均损失填充位置会稀释掉梯度信号而且在类别极不均衡的时候会明显拉偏模型。我在这上面吃过亏一开始指标看起来还行加上 mask 之后同一个模型的 F1 直接涨了好几个点。5.3 用权重形状反推门控顺序如果你需要自己实现一个等价的前向、或者想确认权重到底存在哪可以直接看参数形状gru nn.GRU(10, 20) print(gru.weight_ih_l0.shape) # torch.Size([60, 10]) 3*H_out x input_size print(gru.weight_hh_l0.shape) # torch.Size([20*3, 20]) r, z, n gru.weight_ih_l0.chunk(3, dim0) # 顺序重置门、更新门、新门形状是3 * hidden_size对应三个门在输出方向上拼接。顺序是重置门、更新门、新候选不是直觉上的更新门在前。多层的话权重名会带层号后缀比如weight_ih_l1双向则分_reverse后缀比如weight_ih_l0_reverse。写自定义初始化的时候按这些名字去遍历比手动枚举稳妥得多。顺带说一句参数初始化PyTorch 默认用均匀分布初始化 GRU 权重范围由hidden_size决定。在序列较长、层数较深的时候默认初始化有时会让前几个 epoch 的梯度偏小我习惯把自己的 embedding 用正态分布初始化GRU 部分保持默认这样就够用了。6. 训练循环里的三件小事隐状态传递、detach 和 dropout6.1 跨 batch 传隐状态必须 detach做语言建模或者需要跨越 batch 边界延续状态的场景时会把上一批的h_n当作下一批的h_0。这时如果不.detach()计算图会一直往后延伸跑不了几步就会抛出 Trying to backward through the graph a second time 的错误。h None for xb, yb in loader: out, h model(xb, h) loss criterion(out, yb) loss.backward() optimizer.step() optimizer.zero_grad(set_to_noneTrue) h h.detach() # 关键一步截断计算图set_to_noneTrue是个小细节比清成零稍微省一点显存和带宽在长序列任务里积少成多。6.2 dropout 只在层与层之间生效nn.GRU的dropout参数不是对输入或输出做丢弃而是对每一层之间的输出做丢弃。所以当num_layers1时这个参数实际上没有任何作用PyTorch 会给出一个提醒。要让 dropout 真正生效至少得两层。如果你的模型是单层 GRU又想加正则化正确的位置是在 embedding 之后、或者在线性分类头之前手动加nn.Dropout。6.3 报错速查把前面提到的坑集中成一张表方便排查时对号入座现象或报错大概率原因处理方式维度能跑通但效果异常差seq 与 batch 维写反或未 pack打印 shape 核对改用 packExpected hidden size 报错h_0 形状与 num_layers/direction 不匹配用num_layers * num_directions计算Trying to backward through the graph a second time跨 batch 传递 h_n 未 detach在传回前h.detach()dropout 参数似乎无效num_layers 为 1增加层数或手动加 DropoutGPU 上性能明显低于预期隐状态非连续或未走高效路径.contiguous()检查 pack 是否生效输入张量形状正确但结果随机反向那半段被误用为句子表示改用h[0]/h[1]拼接或池化我个人固化下来的习惯是任何涉及 GRU 的改动先在 CPU 上用torch.randn造一批小张量把input、output、h_n三个形状打印出来确认无误再切到真实数据上跑。这一步大概花三十秒但能省掉后面几小时盯着 loss 曲线发呆的时间。另外output和h_n的关系建议你也亲手跑一遍验证代码比记十条笔记都管用——我自己就是被打脸之后才真正记住双向配置下反向终态落在时间步 0这件事的。