★ FEATURED ARTICLE 很多人在学 LSTM 的时候都会在一个看似“不是重点”的地方卡很久LSTM 的输出到底长什么样、维度是多少。我印象很深的一次是在给团队画一份 LSTM 输出维度图示画了半个多小时满纸都是方块和箭头结果没人看懂我自己也解释不清。后来我把图彻底重画只保留了“批次轴、时间轴、特征轴”这三样东西的尺寸变化问题一下子就清楚了。这篇文章想讲的就是这张被重新画过的图把 LSTM 一次前向计算里到底产生了哪些张量、每个张量的维度是怎么来的、不同框架下为什么输出形状不一样以及遇到 shape 报错时怎么靠一张维度图快速定位彻底讲清楚。适合刚接触循环神经网络、经常在维度上翻车的读者也适合需要给同事或同学讲明白这个问题的朋友。看完你会发现所谓“输出维度”其实并没有那么玄。1. 先把三类“输出”拆开时间步输出、隐藏状态、单元状态1.1 一次前向计算LSTM 手里同时握着哪几个张量LSTM 并不是把输入向量直接变成输出向量那么简单。它在处理序列的每一个时间步都会同时维护、更新并对外暴露几种不同的数据。如果一上来就盯着“输出”两个字很容易把它们混成一团。假设我们有一个 batch 的样本每条样本在时间步 t 的输入是 x_t形状为 (batch_size, input_dim)。与此同时上一个时间步会传来两个内部状态h_{t-1}形状为 (batch_size, hidden_size)叫做隐藏状态c_{t-1}形状为 (batch_size, hidden_size)叫做单元状态也叫记忆单元。LSTM 内部通过遗忘门、输入门、输出门、候选记忆值来更新这两个状态。遗忘门决定旧记忆保留多少输入门决定新信息写入多少两者组合后得到新的 c_t输出门再决定从 c_t 中读取多少内容从而得到新的 h_t。所以严格来说一个时间步结束后LSTM 至少同时“产出”了三个量当前隐藏状态 h_t形状 (batch_size, hidden_size)当前单元状态 c_t形状 (batch_size, hidden_size)当前时间步的对外输出 y_t在大多数框架里 y_t 就等于 h_t所以形状同样是 (batch_size, hidden_size)。这里最关键的一句话是在 LSTM 里“输出维度”和普通全连接层的“输出维度”含义不完全一样。它不是单一数字而是一组沿时间轴展开的二维张量。1.2 到底哪个才是“输出”取决于你后面接什么很多初学着会把 h_t 和 c_t 搞混尤其是看到 Keras 里return_stateTrue返回三个值的时候完全不知道哪个该用、哪个不该用。其实判断方法很简单看你的下游任务需要什么。如果做的是词级别任务比如序列标注、机器翻译的逐词解码你需要每个时间步都有输出于是把 y_1 到 y_T 全部收集起来叠成一个 (batch_size, seq_len, hidden_size) 的张量。如果做的是句级别任务比如情感分类、文本分类你只需要整个序列编码完毕之后的最后一个隐藏状态 h_T也就是 (batch_size, hidden_size)。这个向量相当于把整句话的信息浓缩成了一个固定长度的向量。而 c_t 呢它通常不会直接交给外面的层主要是在 LSTM 内部以及层与层之间传递记忆用的。不过当你要拿 LSTM 状态去初始化另一个 LSTM 时c_t 也和 h_t 一样重要。用一个表格概括一下输出量形状典型使用场景单个时间步输出 y_t / h_t(batch_size, hidden_size)逐词预测、最后一时间步特征全部时间步输出 y_1...y_T(batch_size, seq_len, hidden_size)序列标注、注意力机制源端表示单元状态 c_t(batch_size, hidden_size)跨层传递记忆、初始化解码器记不住的话就记住一句话对外看 h_t对内传 c_t。1.3 决定输出维度的四个超参数LSTM 的输出维度不是凭空来的它由四个核心参数共同决定。搞懂每个参数管哪一维画图时才不会标错。参数名含义主要影响输出结构的位置batch_size一次并行处理的样本数量输出的第 0 维seq_len时间步数即序列长度输出全部时间步时占据第 1 维input_dim每个时间步输入的特征维度只影响输入不决定隐藏层输出维度hidden_sizeLSTM 隐藏单元数量输出的最后一维也是最核心的一维特别提醒一下input_dim 和 hidden_size 经常被混为一谈。输入特征维度可能是 128隐藏单元数量可能是 64两者可以完全不同。LSTM 内部的连接权重会把任意 input_dim 的输入投影到 hidden_size 上。也就是说只要 hidden_size 定了内部所有状态的最后一维就都是 hidden_size和输入维度是多少没有关系。理解了这四个参数接下来就可以看它们如何随时间轴流动。2. 从输入张量到最终输出维度逐层演变的完整路径2.1 先看输入三维张量 (batch_size, seq_len, input_dim)用一个特别具体的例子来走一遍。假设我们做一个文本分类任务一次处理 32 条句子每条句子固定切成 50 个 token每个 token 用一个 128 维向量表示比如经过嵌入层之后得到了 128 维。那么整个网络的输入张量就是第 0 维32代表批次里 32 条样本第 1 维50代表每条样本有 50 个时间步第 2 维128代表每个时间步的特征维度。也就是 (32, 50, 128)。这里可以把整个输入理解成一摞书每一本是 32 条样本的一页每页是 128 维特征一共 50 页LSTM 一页一页地翻。LSTM 在每一个时间步 t其实只“看见”了第 t 个时间步的数据也就是一个切片 x_t形状是 (batch_size, input_dim)在这个例子里是 (32, 128)。它会拿着这个切片和上一个时间步传来的 h_{t-1}、c_{t-1} 一起计算然后产出新的 h_t 和 c_t。2.2 单个时间步内的维度变化为什么中间所有张量形状都一样继续用上面的例子假设我们设定的 hidden_size 64。在时间步 t下面这些东西的形状分别是输入切片 x_t(32, 128)上一个隐藏状态 h_{t-1}(32, 64)上一个单元状态 c_{t-1}(32, 64)遗忘门、输入门、输出门各自的输出(32, 64)候选记忆值(32, 64)更新后的单元状态 c_t(32, 64)更新后的隐藏状态 h_t(32, 64)。为什么会这样关键在权重矩阵的设计上。以遗忘门为例它的计算方式是f_t σ(W_f · [x_t, h_{t-1}] b_f)这里 [x_t, h_{t-1}] 是把输入和状态拼起来维度是 (32, 128 64) (32, 192)。W_f 的形状是 (hidden_size, input_dim hidden_size)也就是 (64, 192)。两者相乘之后输出维度自然就是 (32, 64)。也就是说输入特征再宽进门之后都会被压缩到 hidden_size 这个尺度上。这也是为什么前面强调 input_dim 不影响输出维度——它影响的只是 W_f 这个矩阵的宽度不是最终表示的长度。如果画成时间步内部的前向计算公式可以写成下面这样i_t σ(W_i · [x_t, h_{t-1}] b_i) f_t σ(W_f · [x_t, h_{t-1}] b_f) o_t σ(W_o · [x_t, h_{t-1}] b_o) c~_t tanh(W_c · [x_t, h_{t-1}] b_c) c_t f_t ⊙ c_{t-1} i_t ⊙ c~_t h_t o_t ⊙ tanh(c_t)把形状一一对应上去你会发现它们的最后一维全部是 hidden_size。2.3 整个序列处理完成后的最终输出现在序列走完 50 个时间步LSTM 的“最终输出”会根据你的后续用途有两种形态。第一种只要最后一个时间步的状态。比如做文本分类你把最后一个时间步的 h_T 拿出来形状是 (batch_size, hidden_size)也就是 (32, 64)。这个向量已经包含了整个序列的压缩信息。第二种把每个时间步的输出都保留。比如做序列标注你需要 y_1 到 y_T 每一个时间步的结果于是把它们沿着时间轴叠起来得到一个三维张量形状是 (batch_size, seq_len, hidden_size)也就是 (32, 50, 64)。注意这两种形态在维度图上必须分清楚。很多框架里return_sequencesFalse和return_sequencesTrue这两个参数控制的就是这个区别。返回最后一步时是二维张量返回全部时间步时是三维张量这两个张量直接决定了你后面全连接层的输入形状。3. 把维度画成人能看懂的样子图示方法与标注规范3.1 常见的 LSTM 结构图为什么解决不了维度问题网上能找到的 LSTM 结构图九成都是画单元内部的“门控电路”sigmoid 方块、tanh 方块、乘号、加号、各种连线。这种图对理解“LSTM 是怎么计算”非常有帮助但对理解“LSTM 输入输出张量长什么样”基本没用。原因很简单结构图把所有连线都画在二维平面上却从来不标注每条线代表的 shape。你不知道那条线是 (batch, hidden) 还是 (seq_len, hidden)更不知道时间轴在哪里展开。所以画“LSTM 输出维度图示”的时候必须换一种思路不画内部电路只画张量与张量之间的形状关系。这不是同一种图前者叫计算图后者叫数据流图。你要用的是后者。3.2 一套可复制的“三轴标注法”我自己比较推荐的做法是画一个大长方体代表输入然后在时间轴上切薄片每一片代表一个时间步。具体分四步走。第一步画出输入长方体三个维度分别对应 batch、seq_len、input_dim。比如 (32, 50, 128)就在图上标明这三条轴各自代表什么。注意不是画三个数字就完了而是要让看图的人一眼就能明白“时间在竖向走样本在横向铺开特征在纵深方向”。第二步把输入长方体沿着时间轴切成 50 片。每一片是一个 (32, 128) 的截面这就是单个时间步的输入。这样做的好处是把“三维输入”和“每个时间步实际看到的东西”直观连接了起来。第三步在每一个切片后面画一个 LSTM 单元块单元块旁边标注 hidden_size 64。然后从单元块引出两条平行输出线一条是 h_t一条是 c_t两条线都标注为 (32, 64)。h_t 除了输出给外部还要连向下一个时间步c_t 主要用虚线连向下一个时间步表示内部传递。第四步在图的右侧画两个候选输出区域如果把所有 h_t 纵向叠起来就是一个 (32, 50, 64) 的长方体如果只取最后一个 h_T就是一个 (32, 64) 的平面块。这两条路径要用不同颜色的框区分开并且在旁边写清楚“路径 A全序列输出”“路径 B最后状态输出”。这套方法不需要画得多精细关键就是每个张量旁边必须写完整的 shape 数字而不是只写“hidden state”“output”这种名称。3.3 用 Python 快速生成维度示意如果你需要把图放进文档或者做展示可以直接用 matplotlib 画一个简单的维度示意。下面这段代码可以生成最基本的输入、LSTM、两种输出形态import matplotlib.pyplot as plt fig, ax plt.subplots(figsize(9, 5)) def draw_block(x, y, w, h, label, color#dbe9f6): ax.add_patch(plt.Rectangle((x, y), w, h, facecolorcolor, edgecolor#333333, lw1.5)) ax.text(x w / 2, y h / 2, label, hacenter, vacenter, fontsize12) # 输入张量 draw_block(0, 2, 2.6, 1.4, 输入\n(32, 50, 128), color#dbe9f6) # LSTM 单元 draw_block(3.4, 2.2, 1.0, 1.0, LSTM\nhidden64, color#ffe8cc) # 最后状态输出 draw_block(5.4, 3.0, 1.8, 0.6, h_T\n(32, 64), color#d5f5e3) # 全序列输出 draw_block(5.4, 0.7, 1.8, 1.4, 全序列输出\n(32, 50, 64), color#fde2e2) ax.set_xlim(-0.5, 7.8) ax.set_ylim(-0.2, 4.2) ax.axis(off) plt.title(LSTM 输出维度示意图) plt.show()跑完这段代码你会得到一个非常朴素的示意输入框、LSTM 框、两个不同形状的输出框。实际使用中你可以再加箭头、加批次轴说明。画图工具不重要重要的是每一张框图上都必须有真实的 shape 标注。3.4 画图时最容易犯的三个错误第一把 seq_len 当成 hidden_size。Seq_len 是时间轴长度hidden_size 是特征数量两者完全不同。我在图上见过有人在 LSTM 输出框里写 (32, 50)这要么是没搞清楚时间轴要么是没搞清楚输出到底取没取全序列。第二把 c_t 和 h_t 画成不同维度。标准实现里两者最后一维都是 hidden_size多数框架甚至会把它们打包成同样形状的两个张量。如果图上画成 c_t 比 h_t 多一维或少一维后面接代码时会误导自己。第三只画结构图而不标 shape。画了 200 个圆圈别人还是不知道实际张量怎么流。记住这份图的唯一目的是表达形状关系不是表达数学计算。4. 跨框架对比Keras 与 PyTorch 输出维度的真实差异4.1 Keras 中 return_sequences 与 return_state 的使用Keras 的 LSTM 层有三个关键参数return_sequences、return_state、以及比较少用的time_major。return_sequencesFalse时层输出形状是 (batch_size, units)。这里的 units 就是我们说的 hidden_size。return_sequencesTrue时层输出形状是 (batch_size, seq_len, units)。因为要保留每个时间步的输出。return_stateTrue时返回值会变成一个列表或者一个元组包含三个元素output、h、c。即使return_sequencesFalse三个返回值也都在别只接一个变量导致报错。看一个最简单的 Keras 例子from tensorflow.keras import layers x layers.Input(shape(50, 128)) # 不考虑 batch 维 lstm layers.LSTM(64, return_sequencesTrue, return_stateTrue) out, h, c lstm(x) print(out.shape) # (None, 50, 64) print(h.shape) # (None, 64) print(c.shape) # (None, 64)注意这里第一个返回值 out 在return_sequencesTrue时是完整的三维张量。如果你只想要最后一个时间步可以直接用out[:, -1, :]。很多初学者以为第一个返回值就是 h 状态其实不是。第一个返回是“层输出”第二个才是隐藏状态 h。只有在return_sequencesFalse的时候第一个返回值才会退化成 (None, 64)看起来像 h但它是 h_T 的另一种表达方式。4.2 PyTorch 的 batch_first 与 h_n、c_nPyTorch 里的 LSTM 用起来比 Keras 更底层维度也更直白但同时更容易把索引写错。默认情况下PyTorch LSTM 的输入是 (seq_len, batch_size, input_dim)也就是时间轴在第一维。如果你用惯了 Keras 那种“batch 在前”的写法第一次跑 PyTorch 十有八九会报维度错误。解决办法是设置batch_firstTrue。看代码import torch import torch.nn as nn lstm nn.LSTM(input_size128, hidden_size64, num_layers1, batch_firstTrue) x torch.randn(32, 50, 128) # (batch, seq, feature) out, (h_n, c_n) lstm(x) print(out.shape) # (32, 50, 64) print(h_n.shape) # (1, 32, 64) print(c_n.shape) # (1, 32, 64) # 最后一个时间步的输出 last_output out[:, -1, :] # (32, 64)这里的 h_n 和 c_n 都是三维张量第 0 维代表层数。单层单向 LSTM 时第 0 维就是 1。如果设置了batch_firstFalse输入变成 (50, 32, 128)out 变成 (50, 32, 64)取最后一个时间步就要用out[-1, :, :]很多人都是在这里把索引方向搞反的。还有一个很容易忽略的细节out[:, -1, :]与h_n[-1]是不是完全相等在单层单向 LSTM 里两者的确相等。但在多层 LSTM 里out 永远是最后一层的输出序列h_n[-1] 也是最后一层的最终隐藏状态这两个仍然对应。不过如果你想取第一层的中间状态就必须通过 h_n[0] 来拿out 里面是没有的。4.3 双向 LSTM 与多层 LSTM 的维度翻倍规则双向 LSTM 会让输出维度直接翻倍因为每一时间步会把正向隐藏状态和反向隐藏状态拼起来。假设 hidden_size 64双向之后输出最后一维变成 128。Keras 里这样写from tensorflow.keras import layers x layers.Input(shape(50, 128)) lstm layers.Bidirectional( layers.LSTM(64, return_sequencesTrue) )(x) print(lstm.shape) # (None, 50, 128)PyTorch 里则是import torch.nn as nn lstm nn.LSTM(input_size128, hidden_size64, num_layers2, batch_firstTrue, bidirectionalTrue) x torch.randn(32, 50, 128) out, (h_n, c_n) lstm(x) print(out.shape) # (32, 50, 128) print(h_n.shape) # (4, 32, 64)这里的 h_n 第 0 维是 4因为 2 层乘以 2 个方向。不要以为它是 batch 或 seq 维。取值的时候如果只要最后一个时间步用于分类我习惯直接用out[:, -1, :]这样不用去管方向拼接的顺序拿到的一定是拼接后的完整 128 维状态。把常见配置整理成下面这张对照表画图时直接参考配置Keras 输出PyTorch 输出单层单向只要最后状态(None, 64)out[:, -1, :] 为 (batch, 64)单层单向要全部时间步(None, 50, 64)out 为 (batch, 50, 64)双层单向只要最后状态(None, 64)h_n[-1] 为 (batch, 64)单层双向要全部时间步(None, 50, 128)out 为 (batch, 50, 128)双层双向只要最后状态(None, 128)out[:, -1, :] 为 (batch, 128)5. 维度图在排查实战里的用法踩过的坑与自查清单5.1 三个高频维度事故的重现与定位事故一全序列输出直接接全连接层。模型想对整句话做情感分类LSTM 设置成了return_sequencesTrue得到 (32, 50, 64)然后直接接Dense(1)。你以为会得到一个二分类 logits实际得到的是 (32, 50, 1)每一个时间步一个预测。这时维度图能帮你一眼看到问题你画的是整条输出路径但分类任务需要走的是 h_T 那条路径。要么改成return_sequencesFalse要么把out[:, -1, :]或时间轴池化之后再接全连接层。事故二编码器-解码器的状态初始化错误。我用多层 LSTM 做编码器取 h_n 想初始化解码器。当时 h_n 的形状是 (num_layers, batch, hidden)我却在代码里传成了 (batch, hidden)结果解码器第一个时间步直接维度崩溃。事后对着维度图才发现h_n 的第 0 维是层数不是批次。图里如果早一点把 h_n 的三个维度分别标上“层数、批次、隐藏维”就不会踩这个坑。事故三注意力打分时维度错乘。在实现 attention 时查询向量来自解码器的当前隐藏状态 h_t形状是 (batch, hidden)键值向量来自编码器的全部输出形状是 (batch, seq_len, hidden)。打分计算时需要把 h_t 扩展成 (batch, 1, hidden)再和 (batch, seq_len, hidden) 做点积得到一个 (batch, seq_len, 1) 的注意力分数。少一个维度扩展就会得到一堆奇怪的 broadcasting 结果。这种问题靠肉眼看代码很难发现但只要你把每个参与计算的张量形状画到图上错配位置通常三秒就暴露了。5.2 维度自查清单每次写完一个含 LSTM 的模型我会在脑子里过一遍下面这个清单输入张量的维度顺序是否和框架默认格式一致PyTorch 默认 seq 在前Keras 默认 batch 在前到底需要全序列输出还是最后一个时间步输出代码和维度图选的是同一条路径吗双向设置之后输出最后一维是否已经乘以 2多层设置之后h_n 的第 0 维是层数还是层数乘以方向数用了return_stateTrue时返回元组里的顺序是否确认过output、h、c 有没有接反图里填写的 hidden_size 和代码里的 units / hidden_size 保持一致了吗每个关键节点之后是否有临时打印 shape 验证而不是习惯性往下写。这些问题只要有一个没确认模型训练时大概率会在某些奇怪的地方报错而且是那种看起来像数据问题、实际是维度问题导致的报错。5.3 两个实测下来很好用的小习惯第一个习惯写在代码注释里。我在声明 LSTM 的上一行和下一行一定会写上输入输出 shape 的注释比如# 输入: (batch32, seq50, feature128) lstm_out, h_n, c_n lstm(x) # 输出: (32, 50, 64), h_n(1, 32, 64), c_n(1, 32, 64)这个习惯看起来笨但换项目、换框架、换同事的时候真的能省下大量沟通成本。第二个习惯画图时把“批次维”统一放在最左边“时间维”统一放在中间“特征维”统一放在最右边。不要这次画 (batch, seq, feature)下次画 (seq, batch, feature)否则对照代码时很容易产生错觉。我自己在项目里见过太多因为维度混乱导致的返工而每一次最后定位到根因时原因都非常朴素不知道某个中间张量到底是什么形状。所以我才一直推荐把维度图作为学习 LSTM 的第一张图而不是最后一张图。把输入输出形状先刻在脑子里再去研究门控原理整个模型的抽象感会少很多。
阅读完成 · 觉得有帮助？
