1. 從一次張量求和錯(cuò)誤說起為什么需要理解cumsum的維度最近在調(diào)試一個(gè)序列預(yù)測模型的損失函數(shù)時(shí)遇到了一個(gè)讓我排查了半天的“詭異”問題。我的目標(biāo)是計(jì)算一個(gè)批次batch中每個(gè)樣本sample在時(shí)間步timestep上的累積損失。直覺上我寫下了類似torch.cumsum(loss_per_timestep, dim1)的代碼心想這應(yīng)該沿著時(shí)間維度累加。然而最終得到的張量形狀和數(shù)值完全不符合預(yù)期導(dǎo)致后續(xù)的梯度計(jì)算和模型更新出現(xiàn)了難以察覺的偏差。問題的根源恰恰在于我對torch.cumsum累積求和函數(shù)中dim參數(shù)的理解不夠透徹。在PyTorch、NumPy等張量計(jì)算庫中cumsum是一個(gè)看似簡單但維度邏輯極易混淆的操作。它不像sum(dim1)那樣在指定維度上求和后該維度就消失了降維而是會保留原維度但該維度上的每個(gè)元素都變成了從起始位置到當(dāng)前位置的累積和。這個(gè)“保留原維度”的特性加上PyTorch張量可能擁有的批量batch、通道channel、高度height、寬度width等多維結(jié)構(gòu)使得dim參數(shù)的選擇變得至關(guān)重要。選錯(cuò)維度不僅會得到錯(cuò)誤的數(shù)值結(jié)果更可怕的是有時(shí)形狀shape可能看起來“合理”但數(shù)據(jù)的內(nèi)在邏輯已經(jīng)全亂這種靜默的錯(cuò)誤在深度學(xué)習(xí)訓(xùn)練中極具破壞性。因此今天我們就來徹底剖析torch.cumsum把它的維度邏輯掰開揉碎講清楚。無論你是剛接觸PyTorch的新手還是想鞏固基礎(chǔ)的老手理解了這個(gè)函數(shù)你就能避免很多因維度誤解導(dǎo)致的坑。2.torch.cumsum的核心機(jī)制在保留中累積在深入維度之前我們必須先建立對torch.cumsum操作本質(zhì)的清晰認(rèn)知。它與普通的sum有根本性的區(qū)別。torch.sum(tensor, dimd) 這是一個(gè)歸約操作。它在指定的維度d上將所有元素相加最終得到一個(gè)單一的值。這個(gè)操作會消除消除維度d。例如一個(gè)形狀為[2, 3, 4]的張量在dim1上求和后形狀變?yōu)閇2, 4]中間的維度3消失了。torch.cumsum(tensor, dimd) 這是一個(gè)掃描操作。它同樣沿著指定的維度d進(jìn)行計(jì)算但不是一次性歸約而是進(jìn)行前綴和掃描。對于維度d上的第i個(gè)位置其值被替換為從該維度第0個(gè)位置到第i個(gè)位置所有元素的累加和。關(guān)鍵點(diǎn)在于輸出張量的形狀與輸入張量完全一致維度d被保留了下來只是其上的數(shù)據(jù)變成了累積值。讓我們用一個(gè)極其簡單的例子來可視化這個(gè)過程。假設(shè)我們有一個(gè)一維張量import torch a torch.tensor([1, 2, 3, 4])torch.cumsum(a, dim0)的計(jì)算過程如下輸出位置0a[0] 1輸出位置1a[0] a[1] 1 2 3輸出位置2a[0] a[1] a[2] 1 2 3 6輸出位置3a[0] a[1] a[2] a[3] 1 2 3 4 10所以結(jié)果是tensor([1, 3, 6, 10])?？梢钥吹叫螤钜廊皇?4,)但每個(gè)位置的值都包含了之前所有位置的信息。注意dim參數(shù)在PyTorch中決定了計(jì)算發(fā)生的“軸”。對于一維張量dim只能是0這很直觀。困惑主要始于二維及更高維張量。3. 二維與高維張量的維度詳解dim如何決定累加方向當(dāng)張量維度升高時(shí)dim參數(shù)就像是指揮官告訴cumsum“請沿著這個(gè)方向一行一行或一列一列地進(jìn)行累積掃描”。理解這一點(diǎn)最好的方式就是看例子。3.1 二維張量矩陣的經(jīng)典場景假設(shè)我們有一個(gè)3行4列的矩陣可以把它想象成一個(gè)迷你數(shù)據(jù)集有3個(gè)樣本每個(gè)樣本有4個(gè)特征。B torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) print(‘B.shape:‘, B.shape) # torch.Size([3, 4])情況一dim0沿行累加即跨樣本累積cumsum_dim0 torch.cumsum(B, dim0) print(cumsum_dim0) # tensor([[ 1, 2, 3, 4], # 第一行不變 # [ 6, 8, 10, 12], # 第二行 第一行 第二行 # [15, 18, 21, 24]]) # 第三行 第一行 第二行 第三行操作解讀dim0表示沿著“行”的方向第0維。你可以想象一列一列地看。以第一列[1, 5, 9]為例輸出行0列01輸出行1列01 5 6輸出行2列01 5 9 15結(jié)果形狀依然是[3, 4]。每一列都獨(dú)立地完成了從上到下的累積。這在某些場景下可以理解為隨著樣本索引的增加特征的累積效應(yīng)。情況二dim1沿列累加即在樣本內(nèi)部跨特征累積cumsum_dim1 torch.cumsum(B, dim1) print(cumsum_dim1) # tensor([[ 1, 3, 6, 10], # 第一行1, 12, 123, 1234 # [ 5, 11, 18, 26], # 第二行5, 56, 567, 5678 # [ 9, 19, 30, 42]]) # 第三行9, 910, 91011, 9101112操作解讀dim1表示沿著“列”的方向第1維。你可以想象一行一行地看。以第一行[1, 2, 3, 4]為例輸出行0列01輸出行0列11 2 3輸出行0列21 2 3 6輸出行0列31 2 3 4 10結(jié)果形狀依然是[3, 4]。每一行都獨(dú)立地完成了從左到右的累積。這是更常見的用法例如計(jì)算一個(gè)序列時(shí)間步的累積概率、累積損失或前綴和。實(shí)操心得對于二維矩陣一個(gè)快速的記憶方法是——dim的值指定了累積發(fā)生的方向而另一個(gè)維度則保持獨(dú)立。dim0是“豎著累”dim1是“橫著累”。在深度學(xué)習(xí)里數(shù)據(jù)通常組織為[batch_size, sequence_len, features]如果你想要計(jì)算每個(gè)樣本內(nèi)部時(shí)間步上的累積和那么dim通常應(yīng)對應(yīng)sequence_len所在的維度很可能是1。3.2 三維張量的實(shí)戰(zhàn)推演三維張量在深度學(xué)習(xí)中無處不在例如批量的圖像[B, C, H, W]或批量的序列數(shù)據(jù)[B, T, D]。這里我們以序列數(shù)據(jù)為例構(gòu)造一個(gè)張量2個(gè)樣本每個(gè)樣本有3個(gè)時(shí)間步每個(gè)時(shí)間步有4個(gè)特征。C torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], [[13,14,15,16], [17,18,19,20], [21,22,23,24]]]) print(‘C.shape:‘, C.shape) # torch.Size([2, 3, 4]) # 可以理解為C[0] 是第一個(gè)樣本的3x4矩陣C[1]是第二個(gè)樣本的3x4矩陣。現(xiàn)在我們分別沿著三個(gè)維度進(jìn)行累積求和。情況一dim0沿批次維度累積cs_dim0 torch.cumsum(C, dim0) print(cs_dim0.shape) # torch.Size([2, 3, 4]) print(cs_dim0) # 第一個(gè)樣本索引0: 和原始C[0]一樣 # tensor([[ 1, 2, 3, 4], # [ 5, 6, 7, 8], # [ 9, 10, 11, 12]]) # 第二個(gè)樣本索引1: 變成了 C[0] C[1] # tensor([[ 113, 214, 315, 416], # [ 517, 618, 719, 820], # [ 921, 1022, 1123, 1224]]) # 即 # tensor([[14, 16, 18, 20], # [22, 24, 26, 28], # [30, 32, 34, 36]])解讀dim0是批次維度。輸出張量的第一個(gè)“切片”cs_dim0[0]就是輸入的第一個(gè)樣本C[0]。輸出的第二個(gè)“切片”cs_dim0[1]是第一個(gè)樣本和第二個(gè)樣本的逐元素和。這相當(dāng)于跨樣本的累積在大多數(shù)訓(xùn)練場景下不常用但在某些特殊的在線學(xué)習(xí)或累積統(tǒng)計(jì)場景可能有用。情況二dim1沿時(shí)間步維度累積cs_dim1 torch.cumsum(C, dim1) print(cs_dim1.shape) # torch.Size([2, 3, 4]) print(‘第一個(gè)樣本的累積和:‘) print(cs_dim1[0]) # 對于C[0] [[1,2,3,4], [5,6,7,8], [9,10,11,12]] # 沿dim1第二個(gè)維度即行/時(shí)間步累積 # 時(shí)間步0: [1,2,3,4] # 時(shí)間步1: [15, 26, 37, 48] [6, 8, 10, 12] # 時(shí)間步2: [159, 2610, 3711, 4812] [15, 18, 21, 24] # 所以 cs_dim1[0] 是 # tensor([[ 1, 2, 3, 4], # [ 6, 8, 10, 12], # [15, 18, 21, 24]])解讀這是序列建模中最常用的模式dim1對應(yīng)的是序列長度/時(shí)間步維度。操作在每個(gè)樣本內(nèi)部獨(dú)立進(jìn)行。對于每個(gè)樣本它計(jì)算每個(gè)特征在所有時(shí)間步上的累積和。輸出張量在dim1上的每個(gè)位置t都包含了從時(shí)間步0到t的所有信息。這對于計(jì)算累積注意力權(quán)重、累積損失、或?qū)崿F(xiàn)像CUMSUM這樣的遞歸層至關(guān)重要。情況三dim2沿特征維度累積cs_dim2 torch.cumsum(C, dim2) print(cs_dim2.shape) # torch.Size([2, 3, 4]) print(‘第一個(gè)樣本第一個(gè)時(shí)間步的累積和:‘) print(cs_dim2[0, 0]) # 對應(yīng) C[0,0] [1,2,3,4] # 輸出tensor([ 1, 3, 6, 10]) # 1, 12, 123, 1234解讀dim2對應(yīng)特征維度。它在每個(gè)樣本的每個(gè)時(shí)間步內(nèi)部對特征進(jìn)行累積。這在某些特定場景下有用例如當(dāng)你有一組按特定順序排列的特征如分位數(shù)并且你想計(jì)算它們的累積分布時(shí)。為了更清晰地對比我們可以用下表總結(jié)三維張量[B, T, D]上不同dim的含義dim參數(shù)累加方向解讀輸出形狀典型應(yīng)用場景dim0跨樣本累積。output[b, t, d] sum_{i0}^ input[i, t, d][B, T, D]較少使用或在在線學(xué)習(xí)中對歷史批次數(shù)據(jù)進(jìn)行累積統(tǒng)計(jì)。dim1樣本內(nèi)跨時(shí)間累積。output[b, t, d] sum_{j0}^{t} input[b, j, d][B, T, D]最常用。序列模型中的累積注意力、累積損失、前綴和計(jì)算、單調(diào)對齊等。dim2時(shí)間步內(nèi)跨特征累積。output[b, t, d] sum_{k0}^bextvpl input[b, t, k][B, T, D]特定特征工程如計(jì)算累積分布函數(shù)CDF、特征重要性累積等。4. 反向累積與cumsum的高級參數(shù)dim與dtypetorch.cumsum的功能不止于此它還有一些有用的參數(shù)可以幫助我們應(yīng)對更復(fù)雜的情況。4.1 反向累積torch.cumsum(..., dimd).flip(dimd)標(biāo)準(zhǔn)的cumsum是從維度的起始位置索引0開始向前累積。但有時(shí)我們需要從末尾開始反向累積。PyTorch沒有直接的cumsum_reverse函數(shù)但可以通過組合操作實(shí)現(xiàn)。需求對于序列[1, 2, 3, 4]想要得到[10, 9, 7, 4]即[1234, 234, 34, 4]。實(shí)現(xiàn)技巧x torch.tensor([1, 2, 3, 4]) # 方法先翻轉(zhuǎn)再正向累積然后再翻轉(zhuǎn)回來 reverse_cumsum torch.cumsum(x.flip(dims(0,)), dim0).flip(dims(0,)) print(reverse_cumsum) # tensor([10, 9, 7, 4])x.flip(dims(0,))將張量翻轉(zhuǎn)得到[4, 3, 2, 1]。torch.cumsum(..., dim0)正向累積得到[4, 7, 9, 10]。.flip(dims(0,))再次翻轉(zhuǎn)得到最終結(jié)果[10, 9, 7, 4]。對于高維張量只需在flip函數(shù)中指定要翻轉(zhuǎn)的維度即可。這個(gè)技巧在計(jì)算雙向RNN的隱藏狀態(tài)或某些需要“未來信息”累積的場景中非常實(shí)用。4.2 數(shù)據(jù)類型控制dtype參數(shù)累積求和可能導(dǎo)致數(shù)值溢出尤其是使用int32或float16類型時(shí)。torch.cumsum允許你指定輸出張量的數(shù)據(jù)類型。x_int torch.tensor([100, 200, 300], dtypetorch.int16) # 直接累積可能溢出int16范圍-32768 ~ 32767 # print(torch.cumsum(x_int, dim0)) # 可能得到錯(cuò)誤結(jié)果 # 指定輸出為更高精度的類型 result torch.cumsum(x_int, dim0, dtypetorch.int32) print(result) # tensor([100, 300, 600], dtypetorch.int32) # 或者直接轉(zhuǎn)換為浮點(diǎn)數(shù) result_float torch.cumsum(x_int.float(), dim0) print(result_float) # tensor([100., 300., 600.])注意事項(xiàng)在處理大數(shù)或長序列累積時(shí)主動使用dtypetorch.float64雙精度或至少torch.float32單精度是避免數(shù)值精度問題的好習(xí)慣。特別是在損失累積或概率累積時(shí)使用float32是標(biāo)準(zhǔn)做法。5. 實(shí)戰(zhàn)場景深度解析cumsum在深度學(xué)習(xí)中的應(yīng)用與避坑指南理解了原理我們來看看cumsum在真實(shí)項(xiàng)目中的用武之地以及那些容易踩進(jìn)去的坑。5.1 場景一序列標(biāo)注任務(wù)中的累積損失計(jì)算假設(shè)我們有一個(gè)批次的中文分詞或命名實(shí)體識別任務(wù)每個(gè)樣本的序列長度不同我們使用了掩碼mask來忽略填充部分padding。損失loss_per_timestep的形狀是[batch_size, max_seq_len]其中被mask的位置損失為0。錯(cuò)誤做法# 假設(shè)我們想計(jì)算每個(gè)樣本的真實(shí)損失各時(shí)間步損失之和 total_loss_per_sample torch.cumsum(loss_per_timestep, dim1)[:, -1] # 錯(cuò)誤這里dim1是在序列長度上累積取最后一個(gè)位置得到的是整個(gè)序列的累積和這看起來沒錯(cuò)。但是如果序列長度不同max_seq_len位置之后全是padding其損失為0所以[:, -1]取到的確實(shí)是所有有效時(shí)間步的損失和。然而問題在于cumsum保留了中間結(jié)果如果后續(xù)操作錯(cuò)誤地使用了整個(gè)cumsum張量而不僅僅是最后一個(gè)值就會引入錯(cuò)誤。更清晰且安全的做法# 方法1直接使用 sum語義更清晰 total_loss_per_sample torch.sum(loss_per_timestep, dim1) # 方法2如果確實(shí)需要中間累積結(jié)果例如用于動態(tài)規(guī)劃再使用 cumsum cumulative_loss torch.cumsum(loss_per_timestep, dim1) # 使用時(shí)明確知道 cumulative_loss[b, t] 代表樣本b到時(shí)間t為止的損失和5.2 場景二注意力機(jī)制中的單調(diào)對齊Monotonic Alignment在語音合成或序列到序列模型中單調(diào)對齊是一種常見約束要求對齊路徑隨時(shí)間單調(diào)前進(jìn)。cumsum可以用來將二元的“是否前進(jìn)”決策0/1轉(zhuǎn)換為累積的“已經(jīng)前進(jìn)的步數(shù)”。# decision: [batch_size, output_seq_len] 值為0停留或1前進(jìn) decision torch.bernoulli(torch.ones(2, 5) * 0.7) # 隨機(jī)生成0/1 print(decision) # tensor([[1., 1., 0., 1., 0.], # [1., 0., 1., 1., 1.]]) # 累積前進(jìn)的步數(shù) alignment_pos torch.cumsum(decision, dim1) print(alignment_pos) # tensor([[1., 2., 2., 3., 3.], # 第3步?jīng)Q策為0所以位置停留在2 # [1., 1., 2., 3., 4.]])這里dim1沿著輸出序列維度累積。alignment_pos的每個(gè)位置指示了到當(dāng)前輸出步為止已經(jīng)消耗了多少個(gè)輸入步。這是實(shí)現(xiàn)單調(diào)注意力Monotonic Attention等機(jī)制的關(guān)鍵一步。5.3 場景三計(jì)算前綴和Prefix Sum以進(jìn)行高效采樣在強(qiáng)化學(xué)習(xí)或自然語言生成中我們經(jīng)常需要根據(jù)一個(gè)概率分布進(jìn)行采樣。如果我們需要從一系列動作或單詞的累積概率中進(jìn)行采樣cumsum就派上用場了。probs torch.tensor([0.1, 0.4, 0.2, 0.3]) # 動作概率和為1 cum_probs torch.cumsum(probs, dim0) print(cum_probs) # tensor([0.1000, 0.5000, 0.7000, 1.0000]) # 生成一個(gè)隨機(jī)數(shù)用于采樣 rand torch.rand(1) # 找到第一個(gè)累積概率大于隨機(jī)數(shù)的索引 selected_action torch.searchsorted(cum_probs, rand).item() print(f“隨機(jī)數(shù) {rand.item():.3f}, 選擇動作 {selected_action}“)這里cum_probs將概率轉(zhuǎn)換成了區(qū)間[0, 1]上的分段界限torch.searchsorted可以高效地找到隨機(jī)數(shù)落入的區(qū)間從而完成采樣。這在處理大批量、多類別的采樣時(shí)非常高效。5.4 常見踩坑點(diǎn)與排查清單維度混淆導(dǎo)致計(jì)算錯(cuò)誤這是最普遍的問題。務(wù)必在操作前打印張量的shape并明確你希望累積發(fā)生在哪個(gè)物理意義上跨樣本、跨時(shí)間、跨特征。黃金法則對形狀為[B, T, D]的張量做cumsum(dim1)你得到的是每個(gè)樣本內(nèi)、每個(gè)特征維度上、隨時(shí)間步的累積。誤用cumsum代替sum如果你只需要最終的總和請使用torch.sum。cumsum會保留所有中間結(jié)果消耗更多內(nèi)存且可能讓后續(xù)代碼邏輯變復(fù)雜。忽略掩碼Mask在序列任務(wù)中如果序列長度不一致通常會有填充和掩碼。對包含填充位置的數(shù)據(jù)進(jìn)行cumsum填充位置通常為0也會被累積進(jìn)去這可能會影響有效位置的值。一種做法是先將填充位置設(shè)為0再進(jìn)行累積。更嚴(yán)謹(jǐn)?shù)淖龇ㄊ鞘褂胻orch.where或構(gòu)造一個(gè)與掩碼相關(guān)的偏移量來修正累積結(jié)果。數(shù)值溢出與精度對整數(shù)或低精度浮點(diǎn)數(shù)進(jìn)行長序列累積很容易溢出。如前所述考慮使用dtype參數(shù)提升計(jì)算精度。梯度傳播cumsum操作是完全可微分的可以無縫融入神經(jīng)網(wǎng)絡(luò)的反向傳播中。每個(gè)輸出位置的梯度都會流回所有參與該位置累積計(jì)算的輸入位置。這在構(gòu)建自定義的可微分層時(shí)非常有用。6. 從cumsum到cumprod與logcumsumexp擴(kuò)展你的累積操作工具箱理解了cumsum它的兩個(gè)“親戚”操作就很容易掌握了。torch.cumprod累積乘積與cumsum邏輯完全一致只是將加法換成乘法。常用于計(jì)算連乘例如計(jì)算一系列條件概率的累積乘積。a torch.tensor([1, 2, 3, 4]) print(torch.cumprod(a, dim0)) # tensor([1, 2, 6, 24]) (1, 1*2, 1*2*3, 1*2*3*4)torch.logcumsumexp數(shù)值穩(wěn)定的對數(shù)空間累積求和這是深度學(xué)習(xí)中一個(gè)極其重要但容易被忽視的函數(shù)。當(dāng)我們需要計(jì)算log(sum(exp(x)))的累積版本時(shí)直接計(jì)算exp很容易導(dǎo)致數(shù)值溢出exp(100)就是一個(gè)巨大的數(shù)。logcumsumexp使用了數(shù)值穩(wěn)定的算法。x torch.tensor([100., 101., 102.]) # 直接計(jì)算 log(sum(exp(x))) 會溢出 # logcumsumexp 可以穩(wěn)定計(jì)算 result torch.logcumsumexp(x, dim0) print(result) # tensor([100.0000, 101.3133, 102.4076]) # 等價(jià)于 log(exp(100)exp(101)), log(exp(100)exp(101)exp(102))但過程穩(wěn)定。它在計(jì)算注意力權(quán)重、歸一化流Normalizing Flows或任何涉及對數(shù)概率求和累積的場景中都是必備工具。7. 性能考量與替代方案何時(shí)該用cumsumtorch.cumsum在GPU上由高度優(yōu)化的CUDA內(nèi)核實(shí)現(xiàn)對于大多數(shù)應(yīng)用來說性能足夠好。然而如果你在處理超長序列比如數(shù)萬甚至更長并且只需要最終的累積和那么使用torch.sum會更快因?yàn)樗苊饬擞?jì)算和存儲所有中間結(jié)果。在自定義CUDA內(nèi)核或使用 Triton 等工具進(jìn)行極致優(yōu)化時(shí)你可能需要考慮更復(fù)雜的并行掃描算法如 Blelloch 掃描來優(yōu)化cumsum的性能。但對于99%的PyTorch應(yīng)用內(nèi)置的torch.cumsum就是最佳選擇。最后分享一個(gè)我個(gè)人的調(diào)試習(xí)慣每當(dāng)我在代碼中寫下cumsum時(shí)我都會立刻寫一個(gè)簡單的小例子用肉眼可驗(yàn)證的數(shù)據(jù)就像本文開頭的例子快速跑一下確認(rèn)dim參數(shù)的效果是否符合我的物理直覺。這個(gè)習(xí)慣幫我省下了無數(shù)小時(shí)的調(diào)試時(shí)間。維度操作是張量編程的核心而cumsum又是其中微妙而強(qiáng)大的一個(gè)花時(shí)間徹底理解它絕對是一筆高回報(bào)的投資。