群里有朋友發(fā)來一張截圖說“LSTM輸出輸出維度怎么看都看不懂”。我猜是“輸入輸出”四個字打重了可問題本身一點都不含糊模型明明能跑但中間張量到底長什么樣很多人完全說不清。LSTM的維度問題幾乎每個入門者都會卡一下因為它的數(shù)據(jù)不是一次性算完而是按時間步滾動更新再加上不同框架的默認排列順序還不一樣input_size、hidden_size、seq_len、batch_size幾個概念攪在一起確實容易繞暈。這篇文章就用表格、簡單示意和幾段能直接復制的小代碼把LSTM的輸入輸出維度一層層拆開。從單層到多層從單向到雙向從多對一到多對多最后把常見的維度報錯也一并解決掉。讀過之后你會發(fā)現(xiàn)LSTM的維度沒那么玄關鍵是把“時間步”“隱藏層大小”“批大小”這三件事對應到張量的具體軸上。1. 先搞清楚三個數(shù)字input_size、hidden_size、seq_lenLSTM的維度問題九成是這三個數(shù)字混在一起造成的。它們都叫“大小”但含義完全不同。1.1 input_size 是“每個時間步的特征數(shù)”不是序列長度input_size指的是單個時間步里攜帶的特征數(shù)量。比如一段文本經(jīng)過詞嵌入后每個詞可能用8維向量表示那input_size8再比如你預測一只股票每個歷史時刻只有“收盤價”一個數(shù)字那input_size1如果有5個傳感器的讀數(shù)那input_size5。很多從全連接網(wǎng)絡轉(zhuǎn)過來的朋友會下意識把input_size理解成“輸入神經(jīng)元個數(shù)”這個類比在LSTM里其實不太準確。LSTM每個時間步都是一個小型計算單元它接收的是一整個向量而不是一個標量。這個向量的長度就是input_size。它跟序列長度完全無關。長度是“有多少個時間步”特征數(shù)是“每個時間步里有多少個數(shù)字”。一個句子有10個詞每個詞用8維向量表示那么seq_len10input_size8。這兩個數(shù)字一旦搞反模型肯定報錯。1.2 hidden_size 是“內(nèi)部記憶的寬度”不是層數(shù)hidden_size表示LSTM單元內(nèi)部狀態(tài)向量的長度你可以把它理解為“這個單元記憶有多寬”。它不是神經(jīng)元的個數(shù)而是每一組狀態(tài)向量的維度。例如hidden_size16意味著在內(nèi)部會產(chǎn)生一個16維的隱狀態(tài)向量h_t和一個16維的細胞狀態(tài)向量c_t。這兩個狀態(tài)向量會在每個時間步不斷被更新用來攜帶上下文信息。hidden_size和input_size不需要相等。哪怕輸入特征只有1維內(nèi)部狀態(tài)也可以設成128維。這就像一個人接收到的外部信息可能很簡單但他腦子里可以反復琢磨、展開很多內(nèi)部細節(jié)。hidden_size越大表達能力越強但計算量和過擬合風險也更大。1.3 seq_len 是“循環(huán)的次數(shù)”不是數(shù)據(jù)維度seq_len是時間步的數(shù)量它決定LSTM要按順序處理多少個輸入。它不體現(xiàn)在單個向量的長度上而是體現(xiàn)在張量的一個軸以及內(nèi)部循環(huán)次數(shù)上。把這三個數(shù)字放回張量里看就清楚多了。一個標準的LSTM輸入是三維張量三個軸分別對應軸0批次大小batch_size軸1序列長度seq_len軸2每個時間步的特征數(shù)input_size所以一個[batch_size, seq_len, input_size]的輸入經(jīng)過LSTM層之后一般會得到[batch_size, seq_len, hidden_size]的輸出。這層變換的本質(zhì)就是“把每個時間步的特征向量從 input_size 映射到 hidden_size”同時借助內(nèi)部狀態(tài)把整個序列的信息累積起來。2. 單層LSTM的維度圖示輸入輸出到底怎么變這一節(jié)是全文的基礎。搞懂單層單向LSTM后多層、雙向、各種變體的維度都是在這個基礎上加的。2.1 最常見輸入形態(tài)是三維張量框架默認順序不一樣LSTM輸入一定是三維張量但同一個數(shù)據(jù)在不同框架里擺放順序不一樣。這是很多新手疑惑的重災區(qū)。框架/參數(shù)輸入張量形狀說明PyTorch 默認[seq_len, batch, input_size]時間軸放在最前面time-majorPyTorch 開啟batch_firstTrue[batch, seq_len, input_size]批量軸放最前更直觀TensorFlow/Keras[batch, seq_len, input_size]默認就是批量軸在前PyTorch 默認不把 batch 放前面是因為LSTM類最初是從論文習慣里的“時間步優(yōu)先”風格繼承過來的訓練時可能在某些并行計算場景下更快。但大多數(shù)業(yè)務代碼里數(shù)據(jù)讀進來通常都是batch在前的所以 PyTorch 也提供了batch_firstTrue來切換。我建議你在寫代碼時統(tǒng)一加batch_firstTrue至少能少一半軸順序方面的糾結(jié)。注意一點這個參數(shù)只影響輸入和輸出張量中“時間軸和批量軸”的位置不影響隱狀態(tài)h_n和細胞狀態(tài)c_n的排列方式。2.2 輸出其實有兩種全部時間步輸出和最終隱狀態(tài)LSTM層會返回兩類東西很多人只盯著其中一類然后就蒙了。第一類是“所有時間步的輸出”通常叫output或者out。在 PyTorch 里它的形狀是batch_firstTrue[batch, seq_len, hidden_size]batch_firstFalse[seq_len, batch, hidden_size]第二類是“最后一個時間步的隱狀態(tài)和細胞狀態(tài)”通常叫h_n和c_n。它們的形狀是[num_layers, batch, hidden_size]注意即使你設置batch_firstTrueh_n的形狀也不會變成[batch, num_layers, hidden_size]。它的第一維永遠是層數(shù)。為什么會不一樣因為out保存的是“每個時間步的最后一層輸出”而h_n保存的是“每一層在最后那個時間步的隱狀態(tài)”。單層單向時out[:, -1, :]和h_n[-1]在數(shù)值上是同一個東西。多層時h_n里面還裝著前面層的最終狀態(tài)out里面則是每一時間步的最后一層輸出兩者不能直接畫等號。2.3 一個可復現(xiàn)的維度打印實驗與其空想不如直接打印維度。下面這段代碼在 PyTorch 里驗證單層單向LSTM的輸入輸出形狀import torch import torch.nn as nn # 輸入特征8維隱狀態(tài)16維batch_firstTrue lstm nn.LSTM(input_size8, hidden_size16, batch_firstTrue) # 4句話每句話5個詞每個詞8維向量 x torch.randn(4, 5, 8) out, (h_n, c_n) lstm(x) print(out.shape:, out.shape) # torch.Size([4, 5, 16]) print(h_n.shape:, h_n.shape) # torch.Size([1, 4, 16]) print(c_n.shape:, c_n.shape) # torch.Size([1, 4, 16]) # 所有時間步輸出里的“最后一步” 和 h_n最后一層數(shù)值一致 print(torch.equal(out[:, -1, :], h_n[-1])) # True在 Keras 里對應的參數(shù)是return_sequences和return_statefrom tensorflow.keras.layers import LSTM # return_sequencesTrue返回所有時間步輸出 [batch, seq_len, hidden_size] # return_stateTrue額外返回最后隱狀態(tài) h 和細胞狀態(tài) c lstm_layer LSTM(16, return_sequencesTrue, return_stateTrue)Keras 的返回值順序容易記混依次是output,h_state,c_state。其中h_state的形狀是[batch, hidden_size]不是[1, batch, hidden_size]。這就跟 PyTorch 不一樣換框架時一定要重新打印一下。3. 維度流動的內(nèi)部邏輯從 input_size 到 hidden_size知道形狀以后還得理解它為什么會這樣流動。否則下次遇到[batch, seq_len, input_size]變成什么還是只能靠猜。3.1 時間步是“循環(huán)”不是廣播LSTM不是像卷積那樣在整個序列上一次性滑完而是嚴格按時間順序循環(huán)處理。理論上LSTM單元在每個時間步只處理一個“長度為input_size”的向量然后更新內(nèi)部狀態(tài)。這個過程可以用一個手動展開的循環(huán)來理解h torch.zeros(batch, hidden_size) c torch.zeros(batch, hidden_size) for t in range(seq_len): x_t x[:, t, :] # 取出第t個時間步的數(shù)據(jù)形狀 [batch, input_size] h, c lstm_cell(x_t, (h, c)) # 內(nèi)部狀態(tài)更新h仍然是 [batch, hidden_size]所以seq_len是循環(huán)的次數(shù)它決定張量中“時間軸”的長度。每個時間步的輸入和輸出向量長度分別由input_size和hidden_size決定。循環(huán)次數(shù)和向量長度是兩回事這也是很多人搞混的根源。3.2 權(quán)重矩陣把外部特征映射到內(nèi)部記憶空間LSTM單元內(nèi)部有一組門控機制輸入門、遺忘門、輸出門、候選記憶每個門本質(zhì)上都是一次線性變換加激活函數(shù)。線性變換的輸入是“當前輸入向量”和“上一個隱狀態(tài)”拼接在一起然后乘上權(quán)重矩陣。忽略掉具體門控細節(jié)后可以簡化理解為input_size 維向量 - 線性變換 - hidden_size 維隱狀態(tài)input_size和hidden_size不相等完全沒問題因為中間有一個hidden_size x (input_size hidden_size)的權(quán)重矩陣在做維度轉(zhuǎn)換。換句話說LSTM把“外部世界的特征表達”映射到“內(nèi)部記憶空間”里。為什么需要這個差異因為外部特征可能很稀疏、很短而內(nèi)部記憶需要足夠?qū)挷拍艹休d長期上下文。比如一個單詞的詞向量只有8維但句子里的語法角色、情感傾向、指代關系都需要被記憶這就不止8維能裝下。3.3 文本分類任務的維度流動全過程用文本分類舉個例子。假設要判斷4句話的情感句子已經(jīng)切好了詞每句話最多5個詞每個詞用8維的詞向量表示。流程大致是這樣原始文本[我 愛 深 度 學 習, 這 部 電 影 好 看, ...] 分詞后變成 token id[4, 5] 經(jīng)過 Embedding(詞表大小, 8) 變?yōu)閇4, 5, 8] 經(jīng)過 LSTM(input_size8, hidden_size16, batch_firstTrue) 變?yōu)閇4, 5, 16] 取最后一個時間步 out[:, -1, :] 變?yōu)閇4, 16] 接 Dense(2) 輸出分類概率[4, 2]這里的[4, 5, 8]就是標準三維輸入。5是時間步數(shù)8是每個詞向量維度4是批量大小。LSTM層輸出[4, 5, 16]每個詞都有一個16維的隱狀態(tài)表示。如果任務是句子級分類只關心整句話的情感那就取最后一個時間步。如果任務是詞性標注需要每個詞都輸出標簽那就直接用[4, 5, 16]接全連接層對每個時間步做分類。3.4 時間序列預測任務的維度流動全過程再舉一個完全不同領域的例子用過去10個時刻的傳感器數(shù)據(jù)預測下一時刻的溫度。原始數(shù)據(jù)是一維的總共10000個時刻。為了喂給LSTM需要做滑動窗口切樣本。假設每個樣本取10個歷史時刻每個時刻只有1個溫度值批量大小取32。此時輸入張量是[32, 10, 1]32是批量大小一次處理32個窗口10是時間步數(shù)即過去的10個時刻1是每個時刻的特征數(shù)即溫度lstm nn.LSTM(input_size1, hidden_size32, batch_firstTrue) out, (h_n, c_n) lstm(x) # out.shape [32, 10, 32] # 預測下一個溫度取序列最后一步再接一個全連接層 last_step out[:, -1, :] # [32, 32] pred fc(last_step) # [32, 1]如果每個時刻有多個傳感器比如溫度、濕度、風速三個值那input_size3輸入形狀就是[32, 10, 3]。LSTM內(nèi)部會自動把每個時間步的3維特征映射到hidden_size維空間。4. 不同任務形態(tài)下的輸出維度變體LSTM的輸入輸出形態(tài)并不只有一種。根據(jù)任務需要大體分為多對一、多對多、一對多三類。維度選擇完全取決于任務目標。4.1 多對一、多對多、一對多的形狀對照表任務形態(tài)典型任務輸入形狀常用輸出多對一文本分類、時間序列點預測[batch, seq_len, input_size][batch, hidden_size]取最后一步多對多同步序列標注、逐幀分類[batch, seq_len, input_size][batch, seq_len, hidden_size]多對多異步機器翻譯、摘要生成編碼器輸入[batch, seq_len_enc, input]解碼器輸入逐個生成解碼器輸出長度不定一對多圖像標題、文本生成條件向量 逐步輸入生成seq_len_out步輸出“一對多”這個詞容易讓人誤以為LSTM原生支持一個輸入、多個輸出。實際它并不是一個LSTM層就能完成的而是靠“把同一個條件向量重復送到每個時間步”或者“把條件向量作為初始隱狀態(tài)”來構(gòu)造的。4.2 多對一任務取“最后一步”的兩個等價方式多對一任務最后都需要把一整條序列壓縮成一個向量。最常用的是取序列最后一個時間步的隱狀態(tài)。在 PyTorch 里有兩種寫法本質(zhì)等價# 方式一從 out 里取最后一步 feature out[:, -1, :] # [batch, hidden_size] # 方式二從 h_n 里取最后一層 feature h_n[-1] # [batch, hidden_size]單層單向時這兩個結(jié)果完全相同。不過有一個小坑如果開了batch_firstTrueout[:, -1, :]是最后一步如果沒開就應該是out[-1, :, :]。所以每次寫之前都要先確認你的輸入軸順序是什么。多層時更需要注意。out是最后一層在每個時間步的輸出h_n是每一層的最終隱狀態(tài)。如果你要的是“最后一層的最終狀態(tài)”h_n[-1]是對的如果你要的是“倒數(shù)第二層的最終狀態(tài)”就得用h_n[-2]。只看out時你拿不到中間層的輸出。4.3 一對多任務靠“重復輸入”和“初始狀態(tài)”實現(xiàn)一對多任務最常見的是文本生成和圖像標題輸入一個條件向量輸出一串序列。LSTM本身沒有“一個輸入產(chǎn)生多個時間步輸出”的固定接口實際操作中有兩種常見構(gòu)造方式。方式一是把條件向量作為LSTM的初始隱狀態(tài)然后在每個時間步把上一時間步的輸出當作當前輸入。這種方式最接近“自回歸生成”的邏輯但訓練時要費心處理“強制教學”和“自由生成”之間的差異。方式二更簡單把同一個條件向量復制seq_len份組成一個[batch, seq_len, condition_dim]的輸入送入LSTM。比如訓練圖像標題時可以把圖像特征向量復制10次輸出10個詞的預測condition img_feature.unsqueeze(1) # [batch, 1, hidden_size] condition_repeat condition.repeat(1, 10, 1) # [batch, 10, hidden_size] out, (h_n, c_n) lstm(condition_repeat) # out: [batch, 10, hidden_size]方式二在訓練時很好寫但生成階段仍然要改成逐步預測因為目標長度不固定。無論哪種方式一對多任務里“輸出序列長度”是你自己決定的LSTM不會自動判斷。4.4 return_sequences 和 return_state 到底怎么配合在 Keras/TensorFlow 里return_sequences和return_state這兩個布爾參數(shù)經(jīng)常把人繞暈。return_sequencesFalse初始默認返回最后一步隱狀態(tài)形狀[batch, hidden_size]return_sequencesTrue返回所有時間步隱狀態(tài)形狀[batch, seq_len, hidden_size]return_stateTrue不管上面那個額外返回最后隱狀態(tài)h和細胞狀態(tài)c在 Keras 中LSTM(return_sequencesTrue, return_stateTrue)的返回值是三元組output, h, c。其中output和h在最后一步是重合的但維度不同output包含完整的時間軸h只代表最后一步。這個細節(jié)非常容易踩坑我見過不少同事把output和h當成一回事結(jié)果后續(xù)全連接層維度怎么都對不上。5. 雙向和多層LSTM的維度疊加單層單向是基礎但實際項目里經(jīng)常要上雙向、上多層。維度這時候會發(fā)生變化但規(guī)則很清晰。5.1 雙向LSTM輸出維度直接翻倍雙向LSTM用兩條獨立的LSTM分別處理原始序列和反向序列然后把它們的輸出拼接起來。所以不管hidden_size設多少最終每個時間步的輸出維度都是2 * hidden_size。在 PyTorch 里lstm nn.LSTM(input_size8, hidden_size16, bidirectionalTrue, batch_firstTrue) out, (h_n, c_n) lstm(x) print(out.shape:, out.shape) # torch.Size([4, 5, 32]) print(h_n.shape:, h_n.shape) # torch.Size([2, 4, 16]) print(c_n.shape:, c_n.shape) # torch.Size([2, 4, 16])out最后一維是32正好是正向16維和反向16維拼起來的結(jié)果。h_n的排列要注意單層雙向時h_n第一維是2第一個是正向最后時間步的隱狀態(tài)第二個是反向最后時間步的隱狀態(tài)。如果你想拿“整個序列的壓縮表示”不要把h_n直接抻平更好的做法是把正向和反向分別取出來拼接h_forward h_n[0] # [batch, 16] h_backward h_n[1] # [batch, 16] h_concat torch.cat([h_forward, h_backward], dim-1) # [batch, 32]Keras 里用Bidirectional(LSTM(...))也是一樣merge_mode默認是concat所以輸出維度是2 * hidden_size。如果不想要翻倍可以把merge_mode設成sum、mul或ave但實際效果要看任務。5.2 多層LSTMh_n按層排布out只表達最后一層多層LSTM也叫堆疊LSTM上一層每個時間步的輸出作為下一層的輸入。因為每一層的hidden_size通常設成相同值所以out的形狀不會隨層數(shù)變化仍然是[batch, seq_len, hidden_size]。變化的是h_n和c_n。它們的第一維會等于num_layers如果雙向則等于num_layers * 2。這意味著每層都有自己獨立的最終隱狀態(tài)互不混雜。lstm nn.LSTM(input_size8, hidden_size16, num_layers2, batch_firstTrue) out, (h_n, c_n) lstm(x) print(out.shape:, out.shape) # torch.Size([4, 5, 16]) print(h_n.shape:, h_n.shape) # torch.Size([2, 4, 16])第一層和第二層的隱狀態(tài)各有各的含義。取最后一層的最后一時間步可以用out[:, -1, :]或者h_n[-1]。取第一層的最后一時間步只能用h_n[0]out里拿不到。如果你在雙層的 h_n 里看到第一維是2別以為那是“方向”它可能是“層”。更復雜的模型里“層”和“方向”同時存在時第一維會是num_layers * num_directions順序一般是按層排每層內(nèi)部再按方向排。想準確知道打印驗證最穩(wěn)妥。5.3 常見維度拼接和池化操作實際項目里拿到LSTM的輸出后往往不會直接接分類器而是會做各種聚合。這些操作的維度變化也很規(guī)律。操作輸入形狀輸出形狀說明取最后一步[batch, seq_len, hidden][batch, hidden]適合分類任務全局平均池化[batch, seq_len, hidden][batch, hidden]對時間軸取平均全局最大池化[batch, seq_len, hidden][batch, hidden]突出最強信號拼接正反向最后狀態(tài)[batch, hidden]兩個[batch, 2*hidden]雙向表示拼接多層最終狀態(tài)[num_layers, batch, hidden][batch, num_layers*hidden]不常見但可用接全連接層之前一定要確保張量已經(jīng)變成二維[batch, feature_dim]。把[batch, seq_len, hidden]直接塞進全連接層是最常見的報錯來源。6. 常見維度報錯與排查技巧實錄最后分享一些實際排查經(jīng)驗。下面這些報錯和信息提示是我見過的最高頻幾類。6.1 五條高頻報錯信息對照報錯提示通常原因解決思路size mismatch/got ... expected ...輸入最后一維不等于input_size檢查數(shù)據(jù)和LSTM構(gòu)造參數(shù)是否一致尤其是Embedding輸出維度和input_sizeExpected hidden[0] size (..., num_layers, ...)傳入初始隱狀態(tài)h0的層數(shù)維度不對確認h0第一維是num_layers或num_layers * num_directionsInput tensor should have 3 dimensions輸入不是三維張量把二維[batch, features]改成[batch, seq_len, features]mat1 and mat2 shapes cannot be multiplied全連接層吃到了三維張量先用out[:, -1, :]或池化降到二維再接全連接Given groups1, weight size ...常見于Conv和LSTM混用時維度理解偏差檢查卷積層和LSTM層之間是否做過正確的維度調(diào)整這些報錯里最坑的是第三種。有人把[batch, features]當成LSTM輸入以為再加一維就完事結(jié)果時間步維度是1模型不報錯但效果很差。還有人以“反正模型能跑就行”的心態(tài)忽略維度警告結(jié)果訓練出的模型毫無泛化能力。6.2 用隨機輸入打印維度的自檢法結(jié)構(gòu)越復雜越應該在接真實數(shù)據(jù)前先用隨機輸入把維度驗證一遍。我自己養(yǎng)成的習慣是寫一個小函數(shù)import torch import torch.nn as nn def check_lstm(input_size8, hidden_size16, seq_len5, batch4, num_layers1, bidirectionalFalse): lstm nn.LSTM( input_sizeinput_size, hidden_sizehidden_size, num_layersnum_layers, bidirectionalbidirectional, batch_firstTrue, ) x torch.randn(batch, seq_len, input_size) out, (h_n, c_n) lstm(x) print(input :, x.shape) print(out :, out.shape) print(h_n :, h_n.shape) print(c_n :, c_n.shape) check_lstm() check_lstm(bidirectionalTrue) check_lstm(num_layers2)用隨機數(shù)做驗證可以先把“結(jié)構(gòu)錯誤”和“數(shù)據(jù)錯誤”拆開。真實數(shù)據(jù)里常有填充、掩碼、變長序列維度問題混在里面會很難判斷。先用隨機數(shù)跑通結(jié)構(gòu)再套真實數(shù)據(jù)能省下大量排查時間。6.3 給新手的維度記憶口訣和自查清單我自己平時喜歡用幾句口訣快速判斷LSTM維度輸入永遠是三維batch、seq_len、input_size。hidden_size管的是輸出向量長度管不了你有多少個時間步。拿“最后一步”用out[:,-1,:]前提是batch_firstTrue。h_n第一維一定是層數(shù)不是批量也不是序列長度。雙向就是hidden_size翻倍多層就是h_n變厚。接全連接前先降到二維。自查的時候按這個順序問一遍“我的數(shù)據(jù)是不是三維三個軸分別是什么batch_first開沒開我想要每個時間步的輸出還是最后一步如果是雙向維度翻倍了嗎如果是多層h_n的層數(shù)對嗎”這套問題過完絕大多數(shù)維度問題都能定位到具體那一行。最后再分享一個實操小習慣。我的代碼里第一版 forward 一定會加一句print(x.shape, out.shape, h_n.shape)等整個模型結(jié)構(gòu)確認無誤再刪掉。別小看這個土辦法LSTM維度問題在紙面上推十遍不如拿一個隨機張量跑一遍打印出來的形狀就是最誠實的答案?？ňS度不用慌把“時間步、hidden_size、batch”三個數(shù)對到具體軸上剩下的都是重復推演。