1. 項(xiàng)目概述為什么我們需要深入理解repeat()在PyTorch的日常開(kāi)發(fā)里處理張量形狀是繞不開(kāi)的基本功。無(wú)論是為了適配模型輸入還是進(jìn)行批量的數(shù)據(jù)擴(kuò)增你總會(huì)遇到一個(gè)需求把一個(gè)張量沿著某些維度“復(fù)制”若干份。這時(shí)候torch.repeat()函數(shù)就會(huì)頻繁地出現(xiàn)在你的代碼中。乍一看它很簡(jiǎn)單不就是復(fù)制嗎但實(shí)際用起來(lái)新手甚至一些有經(jīng)驗(yàn)的開(kāi)發(fā)者都容易在這里栽跟頭——輸出的張量形狀和自己預(yù)想的不一樣或者內(nèi)存占用突然飆升導(dǎo)致程序崩潰。我自己就踩過(guò)這樣的坑。早期寫(xiě)一個(gè)數(shù)據(jù)預(yù)處理管道需要把單個(gè)樣本的特征向量復(fù)制成一個(gè)小批量。我隨手寫(xiě)了個(gè)x.repeat(32, 1)心想這不就是變成[32, feature_dim]嗎結(jié)果程序報(bào)錯(cuò)了或者得到了一個(gè)完全錯(cuò)誤的形狀調(diào)試了半天才發(fā)現(xiàn)是對(duì)repeat()參數(shù)的理解有偏差。這個(gè)函數(shù)的行為和直覺(jué)上的“復(fù)制”有些微妙的區(qū)別它嚴(yán)格遵循著張量廣播Broadcasting背后的擴(kuò)展邏輯而不是簡(jiǎn)單的拼接。所以今天我們就來(lái)徹底拆解torch.Tensor.repeat()這個(gè)函數(shù)。我會(huì)結(jié)合大量的代碼示例和內(nèi)存布局示意圖不僅告訴你它怎么用更重要的是講清楚它為什么這樣設(shè)計(jì)底層發(fā)生了什么以及在實(shí)際項(xiàng)目中如何高效、正確地使用它避免那些常見(jiàn)的性能陷阱和邏輯錯(cuò)誤。無(wú)論你是剛接觸PyTorch還是想鞏固基礎(chǔ)相信這篇深入的分析都能給你帶來(lái)收獲。2.repeat()函數(shù)的核心機(jī)制與參數(shù)解析2.1 函數(shù)簽名與基本行為首先我們看一下官方定義。repeat()是torch.Tensor的一個(gè)方法其簽名如下Tensor.repeat(*sizes) - Tensor這里的*sizes表示一個(gè)可變參數(shù)它接受一系列整數(shù)指明了沿著原始張量的每一個(gè)維度需要重復(fù)的次數(shù)。最關(guān)鍵的一點(diǎn)是sizes參數(shù)的長(zhǎng)度必須大于或等于原始張量的維度數(shù)ndim。這個(gè)函數(shù)的行為可以概括為在新的維度上分配空間并將原始張量的數(shù)據(jù)復(fù)制填充進(jìn)去。它返回的是一個(gè)全新的張量與原始張量不共享內(nèi)存除非重復(fù)次數(shù)為1。我們來(lái)一個(gè)最簡(jiǎn)單的例子建立直觀感受import torch # 定義一個(gè)1維張量 x torch.tensor([1, 2, 3]) print(fOriginal tensor: {x}, shape: {x.shape}) # shape: torch.Size([3]) # 沿著第0維行方向重復(fù)2次 y x.repeat(2) print(fAfter repeat(2): {y}, shape: {y.shape}) # shape: torch.Size([6]) # 輸出: tensor([1, 2, 3, 1, 2, 3]) # 定義一個(gè)2維張量 x_2d torch.tensor([[1, 2], [3, 4]]) print(f\nOriginal 2D tensor shape: {x_2d.shape}) # torch.Size([2, 2]) # 參數(shù) (3, 2) 意味著沿第0維重復(fù)3次沿第1維重復(fù)2次 y_2d x_2d.repeat(3, 2) print(fAfter repeat(3, 2) shape: {y_2d.shape}) # torch.Size([6, 4]) print(y_2d) # 輸出 # tensor([[1, 2, 1, 2], # [3, 4, 3, 4], # [1, 2, 1, 2], # [3, 4, 3, 4], # [1, 2, 1, 2], # [3, 4, 3, 4]])從輸出可以看到repeat(3, 2)生成了一個(gè)6x4的張量。它是先將原始2x2的矩陣在行第0維上復(fù)制3份堆疊起來(lái)形成一個(gè)6x2的中間形態(tài)然后再將這個(gè)中間形態(tài)的每一行在列第1維上復(fù)制2次最終得到結(jié)果。2.2 參數(shù)sizes的長(zhǎng)度規(guī)則與“維度補(bǔ)齊”這是最容易出錯(cuò)的地方。repeat()要求提供的sizes元組長(zhǎng)度不能小于原始張量的維度。如果長(zhǎng)度大于原始維度會(huì)發(fā)生什么PyTorch 的處理方式是在原始張量的前面即更高維或說(shuō)左邊插入新的維度。這些新插入的維度其大小默認(rèn)為1。然后repeat()再根據(jù)你提供的sizes參數(shù)對(duì)所有維度包括新插入的進(jìn)行重復(fù)操作。x torch.tensor([[1, 2, 3], [4, 5, 6]]) print(fOriginal shape: {x.shape}) # torch.Size([2, 3]) # 情況一參數(shù)長(zhǎng)度等于原始維度 (最常見(jiàn)) y1 x.repeat(2, 3) # 沿dim0重復(fù)2次沿dim1重復(fù)3次 print(frepeat(2, 3) shape: {y1.shape}) # torch.Size([4, 9]) # 情況二參數(shù)長(zhǎng)度大于原始維度 # 提供的參數(shù)是 (2, 2, 3)。原始x是2維參數(shù)是3維。 # PyTorch會(huì)在x前面加1個(gè)維度使其變成3維: shape [1, 2, 3] # 然后對(duì)這個(gè)新的3維張量執(zhí)行 repeat(2, 2, 3) # 即新dim0重復(fù)2次原dim0現(xiàn)在是dim1重復(fù)2次原dim1現(xiàn)在是dim2重復(fù)3次 y2 x.repeat(2, 2, 3) print(frepeat(2, 2, 3) shape: {y2.shape}) # torch.Size([2, 4, 9]) print(fy2[0, ...] y2[1, ...] ? {torch.all(y2[0] y2[1])}) # True 因?yàn)樾录拥木S度被復(fù)制了 # 情況三參數(shù)長(zhǎng)度小于原始維度 會(huì)報(bào)錯(cuò) try: y3 x.repeat(2) # x是2維只給了1個(gè)參數(shù) except Exception as e: print(fError: {type(e).__name__}: {e}) # 報(bào)錯(cuò)RuntimeError: Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor注意這個(gè)“維度補(bǔ)齊”的規(guī)則是理解repeat()與類(lèi)似函數(shù)如expand()區(qū)別的核心。它意味著repeat()總是從最高維最左邊的維度開(kāi)始匹配參數(shù)。當(dāng)你寫(xiě)下x.repeat(a, b, c)時(shí)你是在對(duì)一個(gè)新的、可能經(jīng)過(guò)維度擴(kuò)展的張量進(jìn)行操作而不是直接對(duì)應(yīng)原始張量的維度。2.3repeat()與expand()的深度對(duì)比很多人會(huì)混淆repeat()和expand()因?yàn)樗鼈兌寄芨淖儚埩啃螤?。但它們的底層邏輯和?nèi)存影響天差地別。特性torch.repeat()torch.expand()核心操作數(shù)據(jù)復(fù)制。在新的內(nèi)存位置創(chuàng)建數(shù)據(jù)的副本。視圖創(chuàng)建。不復(fù)制數(shù)據(jù)只創(chuàng)建一個(gè)新的“視圖”通過(guò)廣播規(guī)則實(shí)現(xiàn)維度的邏輯擴(kuò)展。內(nèi)存影響顯式占用新內(nèi)存。輸出張量與輸入張量不共享存儲(chǔ)。內(nèi)存增長(zhǎng)倍數(shù)為各維度重復(fù)次數(shù)的乘積。幾乎不占額外內(nèi)存。輸出與輸入共享底層數(shù)據(jù)存儲(chǔ)前提是擴(kuò)展的維度大小為1或-1。參數(shù)要求參數(shù)指定的是重復(fù)次數(shù)必須 1。參數(shù)指定的是目標(biāo)形狀只能將大小為1的維度擴(kuò)展到更大-1表示保持原樣。靈活性可以在任何維度上進(jìn)行任意次數(shù)的復(fù)制。只能對(duì)原始大小為1的維度進(jìn)行擴(kuò)展其他維度大小必須匹配或?yàn)?1。典型用途需要物理上復(fù)制數(shù)據(jù)時(shí)如數(shù)據(jù)增廣、構(gòu)造特定模式的數(shù)據(jù)。需要高效廣播時(shí)如將偏置向量加到批量數(shù)據(jù)上。來(lái)看一個(gè)對(duì)比示例# 定義一個(gè)可以廣播的張量 (第0維大小為1) x torch.randn(1, 3, 224, 224) # 形狀: [1, 3, 224, 224] # 使用 expand: 高效不復(fù)制數(shù)據(jù) batch_size 32 x_expanded x.expand(batch_size, -1, -1, -1) # 目標(biāo)形狀: [32, 3, 224, 224] print(fx_expanded shape: {x_expanded.shape}) print(fDo x and x_expanded share memory? {x.storage().data_ptr() x_expanded.storage().data_ptr()}) # 大概率是True # 使用 repeat: 低效復(fù)制數(shù)據(jù) x_repeated x.repeat(batch_size, 1, 1, 1) # 參數(shù): [32, 1, 1, 1] print(fx_repeated shape: {x_repeated.shape}) print(fDo x and x_repeated share memory? {x.storage().data_ptr() x_repeated.storage().data_ptr()}) # False # 嘗試對(duì)非1維度使用expand會(huì)報(bào)錯(cuò) x2 torch.randn(2, 3, 224, 224) try: x2_expanded x2.expand(32, -1, -1, -1) # 第0維是2不是1無(wú)法直接expand到32 except Exception as e: print(fExpand error: {e}) # 此時(shí)必須先用 unsqueeze 增加一個(gè)維度或者使用 repeat實(shí)操心得在決定用repeat還是expand之前先問(wèn)自己一個(gè)問(wèn)題“我后續(xù)需要修改這個(gè)擴(kuò)展后的張量并希望原始張量保持不變嗎” 如果答案是肯定的或者原始維度大小不為1那么repeat()是更安全的選擇因?yàn)樗鼊?chuàng)建了副本。如果只是為了廣播運(yùn)算且原始維度大小為1那么expand()在內(nèi)存和速度上具有巨大優(yōu)勢(shì)。3. 多維張量repeat()的逐維拆解與內(nèi)存分析理解了基本規(guī)則后我們深入到多維場(chǎng)景并分析其內(nèi)存影響。3.1 高維張量的重復(fù)模式對(duì)于三維及以上的張量這在深度學(xué)習(xí)中非常常見(jiàn)如[Batch, Channel, Height, Width]repeat()的行為遵循同樣的“從左到右”的維度匹配規(guī)則。我們可以把它想象成一個(gè)嵌套循環(huán)最外層的循環(huán)對(duì)應(yīng)sizes的第一個(gè)參數(shù)最內(nèi)層的循環(huán)對(duì)應(yīng)最后一個(gè)參數(shù)。# 模擬一個(gè)批量圖像特征圖 [batch, channel, height, width] feat_map torch.tensor([ # batch1, channel2, height2, width3 [ [[ 1, 2, 3], [ 4, 5, 6]], [[ 7, 8, 9], [10, 11, 12]] ] ]) print(fOriginal feat_map shape: {feat_map.shape}) # torch.Size([1, 2, 2, 3]) # 目標(biāo)批量擴(kuò)展到4通道數(shù)復(fù)制到3高寬各復(fù)制2次 # 參數(shù)順序?qū)?yīng)batch_dim, channel_dim, height_dim, width_dim repeated feat_map.repeat(4, 3, 2, 2) print(fAfter repeat(4,3,2,2) shape: {repeated.shape}) # torch.Size([4, 6, 4, 6]) # 我們來(lái)驗(yàn)證一下其中一個(gè)數(shù)據(jù)塊 # 原始張量中feat_map[0, 0, 0, 0] 1 # 在新的張量中由于batch重復(fù)4次channel重復(fù)3次height重復(fù)2次width重復(fù)2次 # 那么 repeated[0, 0, 0, 0], repeated[0, 0, 0, 1], repeated[0, 0, 1, 0]... 等位置都應(yīng)該是1 # 檢查 repeated[0, 0, 0:2, 0:2] 這個(gè)2x2的塊 print(\nChecking block at [0,0,:,:]:) print(repeated[0, 0, 0:2, 0:2]) # 應(yīng)該輸出 # tensor([[1, 1], # [1, 1]]) # 因?yàn)閣idth方向重復(fù)了2次所以[1,2,3]變成了[1,1,2,2,3,3]這里取前兩個(gè)是1和1。3.2 內(nèi)存占用計(jì)算與性能陷阱這是使用repeat()時(shí)必須警惕的一點(diǎn)。repeat()是物理復(fù)制數(shù)據(jù)其輸出的張量所占用的內(nèi)存是原始張量的 $\prod_{i} sizes[i]$ 倍。這里的 $sizes[i]$ 是每個(gè)維度的重復(fù)次數(shù)。計(jì)算公式輸出張量?jī)?nèi)存 ≈ 原始張量?jī)?nèi)存 * (sizes[0] * sizes[1] * ... * sizes[k-1])假設(shè)你有一個(gè)浮點(diǎn)型張量x形狀為[100, 256, 256]數(shù)據(jù)類(lèi)型是float324字節(jié)。原始內(nèi)存$100 \times 256 \times 256 \times 4 \text{ bytes} \approx 26.2 \text{ MB}$。如果你不小心寫(xiě)了x.repeat(10, 1, 1)想在第0維復(fù)制10份。輸出形狀[1000, 256, 256]。輸出內(nèi)存$1000 \times 256 \times 256 \times 4 \text{ bytes} \approx 262 \text{ MB}$。這看起來(lái)似乎沒(méi)問(wèn)題但如果你本意是想把批次從100擴(kuò)展到1000而原始張量的100是批次大小那么repeat(10,1,1)的邏輯是錯(cuò)誤的它把100當(dāng)成一個(gè)整體復(fù)制了10次得到了1000個(gè)樣本但每個(gè)樣本的“內(nèi)容”是重復(fù)的。更可怕的是下面這種錯(cuò)誤# 假設(shè)我們有一個(gè) batch_size32 的輸入 batch_input torch.randn(32, 3, 224, 224) # 約 32*3*224*224*4/1024/1024 ≈ 18.4 MB # 錯(cuò)誤意圖想把每個(gè)樣本在批次內(nèi)再?gòu)?fù)制一份變成64個(gè)樣本 # 錯(cuò)誤寫(xiě)法 try: wrong_output batch_input.repeat(2, 1, 1, 1) # 形狀: [64, 3, 224, 224] print(fMemory of wrong_output: {wrong_output.element_size() * wrong_output.nelement() / 1024**2:.2f} MB) except Exception as e: print(fError (maybe OOM): {e}) # 內(nèi)存會(huì)翻倍到約36.8MB。但邏輯是錯(cuò)的它把32個(gè)樣本作為一個(gè)整體復(fù)制了。 # 正確做法如果你想要每個(gè)樣本在批次內(nèi)連續(xù)重復(fù)需要先 unsqueeze 增加一個(gè)維度 # 步驟 [32, ...] - [32, 1, ...] - repeat(1, 2, 1, 1) - [32, 2, ...] - view(-1, ...) correct_input batch_input.unsqueeze(1) # [32, 1, 3, 224, 224] temp_output correct_input.repeat(1, 2, 1, 1, 1) # [32, 2, 3, 224, 224] correct_output temp_output.view(-1, 3, 224, 224) # [64, 3, 224, 224] print(fShape of correct_output: {correct_output.shape})注意事項(xiàng)在深度學(xué)習(xí)訓(xùn)練中尤其是在數(shù)據(jù)加載或數(shù)據(jù)增強(qiáng)環(huán)節(jié)使用repeat()時(shí)一定要預(yù)估輸出張量的大小。一個(gè)不小心的repeat操作可能導(dǎo)致顯存溢出OOM尤其是在處理高分辨率圖像或大語(yǔ)言模型時(shí)。在不確定的情況下先用print(x.repeat(...).shape)和print(x.repeat(...).element_size() * x.repeat(...).nelement() / 1024**2)計(jì)算一下輸出形狀和內(nèi)存占用MB是一個(gè)很好的習(xí)慣。4. 典型應(yīng)用場(chǎng)景與實(shí)戰(zhàn)代碼示例repeat()函數(shù)在PyTorch項(xiàng)目中有多種實(shí)用的場(chǎng)景下面我們結(jié)合具體代碼來(lái)看。4.1 場(chǎng)景一構(gòu)造常量張量或模式化張量當(dāng)你需要快速生成一個(gè)具有重復(fù)模式的張量時(shí)repeat()比循環(huán)或列表生成式高效得多。# 示例1構(gòu)造一個(gè)棋盤(pán)格狀的掩碼 (Checkerboard Mask) # 基礎(chǔ)單元是一個(gè) 2x2 的矩陣 [[0, 1], [1, 0]] base torch.tensor([[0, 1], [1, 0]], dtypetorch.float32) height, width 8, 8 # 計(jì)算需要在行和列上重復(fù)多少次 repeat_h height // base.size(0) # 8 / 2 4 repeat_w width // base.size(1) # 8 / 2 4 checkerboard base.repeat(repeat_h, repeat_w) print(8x8 Checkerboard Mask:) print(checkerboard) # 示例2為每個(gè)樣本添加一個(gè)固定的位置編碼向量 batch_size 5 seq_len 10 hidden_dim 8 # 假設(shè)我們有一個(gè)固定的位置編碼向量這里隨機(jī)生成代替 pos_encoding torch.randn(1, seq_len, hidden_dim) # [1, 10, 8] # 將其擴(kuò)展到整個(gè)批次 pos_encoding_batch pos_encoding.repeat(batch_size, 1, 1) # [5, 10, 8] print(f\nPosition encoding for batch: shape {pos_encoding_batch.shape}) print(fIs the encoding same for all samples in batch? {torch.all(pos_encoding_batch[0] pos_encoding_batch[1])})4.2 場(chǎng)景二數(shù)據(jù)擴(kuò)增Data Augmentation的離線生成在某些情況下我們可能需要在訓(xùn)練前就生成一些增強(qiáng)數(shù)據(jù)而不是在數(shù)據(jù)加載器中在線增強(qiáng)。repeat()可以方便地將原始數(shù)據(jù)與變換后的數(shù)據(jù)拼接起來(lái)。# 假設(shè)我們有一小批原始數(shù)據(jù) original_data torch.randn(4, 3, 64, 64) # 4張RGB小圖 print(fOriginal data shape: {original_data.shape}) # 模擬一個(gè)水平翻轉(zhuǎn)操作這里用索引反轉(zhuǎn)簡(jiǎn)化表示 flipped_data original_data.flip(dims[-1]) # 沿最后一維寬度翻轉(zhuǎn) # 將原始數(shù)據(jù)和翻轉(zhuǎn)數(shù)據(jù)合并成一個(gè)大的批次 augmented_data torch.cat([original_data, flipped_data], dim0) print(fAugmented data shape (using cat): {augmented_data.shape}) # [8, 3, 64, 64] # 使用 repeat 的替代方案如果擴(kuò)增策略是簡(jiǎn)單的復(fù)制 # 例如將每張圖片復(fù)制3次模擬多種裁剪這里用復(fù)制代替 replicated_data original_data.repeat(3, 1, 1, 1) # 注意這是完全復(fù)制不是不同的裁剪 print(fReplicated data shape (using repeat): {replicated_data.shape}) # [12, 3, 64, 64] # 然后可以對(duì) replicated_data 的每一份應(yīng)用不同的隨機(jī)裁剪4.3 場(chǎng)景三為廣播運(yùn)算準(zhǔn)備張量雖然expand()更適合廣播但有時(shí)原始張量沒(méi)有大小為1的維度我們又不想改變?cè)紨?shù)據(jù)就可以先用view/unsqueeze增加維度再用repeat。# 一個(gè)常見(jiàn)的例子將一維偏置向量加到二維特征矩陣上 batch_size 3 feature_dim 5 features torch.randn(batch_size, feature_dim) # [3, 5] bias torch.randn(feature_dim) # [5] # 直接相加會(huì)報(bào)錯(cuò)因?yàn)樾螤畈黄ヅ?# result features bias # RuntimeError # 方法1使用 unsqueeze 和 expand (更高效推薦) bias_expanded bias.unsqueeze(0).expand(batch_size, -1) # [1,5] - [3,5] result1 features bias_expanded # 方法2使用 unsqueeze 和 repeat (也能工作但復(fù)制了數(shù)據(jù)) bias_repeated bias.unsqueeze(0).repeat(batch_size, 1) # [1,5] - [3,5] result2 features bias_repeated print(fAre results equal? {torch.allclose(result1, result2)}) # True print(fDo bias_expanded and bias share memory? {bias_expanded.storage().data_ptr() bias.storage().data_ptr()}) # True print(fDo bias_repeated and bias share memory? {bias_repeated.storage().data_ptr() bias.storage().data_ptr()}) # False在這個(gè)例子中expand是更好的選擇因?yàn)樗鼪](méi)有發(fā)生實(shí)際的數(shù)據(jù)復(fù)制。repeat雖然結(jié)果正確但產(chǎn)生了不必要的內(nèi)存開(kāi)銷(xiāo)。5. 常見(jiàn)錯(cuò)誤排查與高級(jí)技巧5.1 錯(cuò)誤排查清單在使用repeat()時(shí)遇到的錯(cuò)誤大多源于對(duì)形狀和參數(shù)的不理解。下面是一個(gè)快速排查表錯(cuò)誤現(xiàn)象可能原因解決方案RuntimeError: Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor提供的sizes參數(shù)個(gè)數(shù)少于張量的維度數(shù)。檢查x.ndim和len(sizes)。確保len(sizes) x.ndim。不足時(shí)可以補(bǔ)1。輸出形狀與預(yù)期不符1. 混淆了repeat參數(shù)是“次數(shù)”而不是“目標(biāo)大小”。2. 忽略了“維度補(bǔ)齊”規(guī)則參數(shù)匹配錯(cuò)了維度。1. 記住output_size[i] input_size[i] * sizes[i]。2. 使用print(x.shape)和print(x.ndim)確認(rèn)維度手動(dòng)推算或畫(huà)圖理解補(bǔ)齊規(guī)則。程序顯存溢出 (OOM)repeat倍數(shù)過(guò)大導(dǎo)致輸出張量?jī)?nèi)存激增。計(jì)算輸出元素總數(shù)prod(output_shape)。檢查是否合理?？紤]使用expand如果條件允許或分塊處理。數(shù)據(jù)重復(fù)模式錯(cuò)誤錯(cuò)誤理解了重復(fù)的軸向。例如想按樣本重復(fù)卻按特征重復(fù)了。使用小張量如torch.tensor([[1,2],[3,4]])和簡(jiǎn)單的repeat參數(shù)如(2,3)先做實(shí)驗(yàn)打印結(jié)果驗(yàn)證模式。梯度計(jì)算錯(cuò)誤或丟失repeat()操作本身支持梯度傳播。問(wèn)題可能出在后續(xù)操作。確保計(jì)算圖連貫。如果需要對(duì)repeat的結(jié)果進(jìn)行原位操作注意使用.clone()或避免使用inplaceTrue的操作。5.2 結(jié)合view、reshape和permute進(jìn)行復(fù)雜變換repeat()經(jīng)常需要和形狀操作函數(shù)搭配使用以實(shí)現(xiàn)復(fù)雜的張量變換。# 目標(biāo)將一個(gè)形狀為 [a, b] 的張量轉(zhuǎn)換成 [a*n, b]其中每一行是原始行的重復(fù)。 # 例如將 [[1,2], [3,4]] 用 repeat(2,1) 得到 [[1,2],[1,2],[3,4],[3,4]]但這不是我們想要的。 # 我們想要的是[[1,2],[1,2],[3,4],[3,4]]但順序是交錯(cuò)的不上面已經(jīng)是了。 # 更復(fù)雜的例子想要得到 [[1,2],[3,4],[1,2],[3,4]]即先整體重復(fù)而不是按行重復(fù)。 x torch.tensor([[1, 2], [3, 4]]) print(fOriginal:\n{x}) # 方法A先增加一個(gè)維度再重復(fù)最后展平 # Step1: [2,2] - [2,1,2] x_a x.unsqueeze(1) # 在dim1處插入維度 print(fAfter unsqueeze(1): shape{x_a.shape}) # Step2: 沿新插入的維度重復(fù)2次 [2,1,2] - [2,2,2] x_a_repeat x_a.repeat(1, 2, 1) print(fAfter repeat(1,2,1): shape{x_a_repeat.shape}) # Step3: 合并前兩個(gè)維度 [2,2,2] - [4,2] x_a_final x_a_repeat.view(-1, x.size(-1)) print(fFinal result (view):\n{x_a_final}) # 方法B直接 repeat 然后 permute (轉(zhuǎn)置) # Step1: 直接重復(fù)整個(gè)張量 [2,2] - [4,4]? 不對(duì)。 # 我們需要的是行重復(fù)所以應(yīng)該是 repeat(2,1) - [4,2] x_b_repeat x.repeat(2, 1) print(f\nDirect repeat(2,1):\n{x_b_repeat}) # 輸出 [[1,2],[3,4],[1,2],[3,4]]正是我們想要的。 # 所以對(duì)于這個(gè)簡(jiǎn)單目標(biāo)直接 repeat(2,1) 即可。 # 但如果我們想要列重復(fù)呢即得到 [[1,2,1,2],[3,4,3,4]] x_col_repeat x.repeat(1, 2) print(fDirect repeat(1,2):\n{x_col_repeat})高級(jí)技巧當(dāng)你的重復(fù)邏輯比較復(fù)雜涉及維度的重新排列時(shí)一個(gè)有效的策略是先用unsqueeze在目標(biāo)位置插入大小為1的維度。這為你提供了重復(fù)的“錨點(diǎn)”。然后使用repeat在這個(gè)新維度上進(jìn)行復(fù)制。最后用view或reshape將張量重塑成最終形狀。如果維度順序不對(duì)可能還需要配合permute進(jìn)行維度換位。5.3repeat與repeat_interleave的辨析PyTorch 中還有一個(gè)函數(shù)叫torch.repeat_interleave()它的行為與repeat()不同也經(jīng)常被混淆。tensor.repeat(*sizes): 按維度整體重復(fù)。參數(shù)指定每個(gè)維度重復(fù)的次數(shù)。torch.repeat_interleave(tensor, repeats, dim): 按維度內(nèi)的元素重復(fù)。repeats可以是一個(gè)整數(shù)所有元素重復(fù)相同次數(shù)也可以是一個(gè)指定每個(gè)元素重復(fù)次數(shù)的張量。x torch.tensor([[1, 2], [3, 4]]) # repeat: 整體重復(fù) print(x.repeat(2, 3):) print(x.repeat(2, 3)) # 整體形狀 [4,6] # 輸出 # [[1, 2, 1, 2, 1, 2], # [3, 4, 3, 4, 3, 4], # [1, 2, 1, 2, 1, 2], # [3, 4, 3, 4, 3, 4]] # repeat_interleave: 元素重復(fù) print(\ntorch.repeat_interleave(x, repeats2, dim0):) print(torch.repeat_interleave(x, repeats2, dim0)) # 沿dim0每一行重復(fù)2次 # 輸出 # [[1, 2], # [1, 2], # [3, 4], # [3, 4]] print(\ntorch.repeat_interleave(x, repeatstorch.tensor([1, 2]), dim1):) print(torch.repeat_interleave(x, repeatstorch.tensor([1, 2]), dim1)) # 沿dim1第0列重復(fù)1次第1列重復(fù)2次 # 輸出 # [[1, 2, 2], # [3, 4, 4]]選擇指南如果你需要像“鋪瓷磚”一樣復(fù)制整個(gè)張量塊用repeat()。如果你需要復(fù)制張量?jī)?nèi)部的特定元素或切片例如將每個(gè)樣本重復(fù)不同的次數(shù)用repeat_interleave()。6. 性能優(yōu)化與替代方案探討雖然repeat()很方便但在性能敏感或內(nèi)存受限的場(chǎng)景下我們需要考慮更優(yōu)的方案。6.1 使用expand替代repeat以節(jié)省內(nèi)存這是最重要的優(yōu)化策略。只要源張量在需要擴(kuò)展的維度上大小為1就優(yōu)先使用expand()。import time # 創(chuàng)建一個(gè)大的張量其中一個(gè)維度為1 large_tensor torch.randn(1, 256, 1024, 1024) # 約 1GB (假設(shè)float32) target_repeats 8 # 測(cè)試 repeat 的內(nèi)存和時(shí)間 start time.time() repeated large_tensor.repeat(target_repeats, 1, 1, 1) time_repeat time.time() - start mem_repeat repeated.element_size() * repeated.nelement() / 1024**3 # GB print(frepeat 耗時(shí): {time_repeat:.3f}s, 內(nèi)存占用: {mem_repeat:.2f} GB) # 測(cè)試 expand 的內(nèi)存和時(shí)間 start time.time() expanded large_tensor.expand(target_repeats, -1, -1, -1) time_expand time.time() - start # expand 不分配新的大內(nèi)存底層存儲(chǔ)指針相同 mem_expand large_tensor.element_size() * large_tensor.nelement() / 1024**3 # 仍然是原始大小 print(fexpand 耗時(shí): {time_expand:.3f}s, ‘視圖’內(nèi)存占用: {mem_expand:.2f} GB (實(shí)際不新增)) print(fexpand 結(jié)果與 repeat 結(jié)果數(shù)值相等嗎 {torch.allclose(expanded, repeated)})6.2 使用torch.cat或torch.stack進(jìn)行可控拼接當(dāng)重復(fù)的份數(shù)不多或者重復(fù)邏輯不是簡(jiǎn)單的整體復(fù)制時(shí)顯式使用cat或stack可能更清晰有時(shí)在反向傳播時(shí)計(jì)算圖也更簡(jiǎn)單。# 假設(shè)我們要將同一個(gè)張量重復(fù)3次形成一個(gè)新的維度 x torch.randn(100, 200) repeat_times 3 # 方法1: repeat result_repeat x.repeat(repeat_times, 1, 1) # 錯(cuò)誤這會(huì)把100也重復(fù)。 # 正確做法是先 unsqueeze result_repeat_correct x.unsqueeze(0).repeat(repeat_times, 1, 1) # [3, 100, 200] # 方法2: stack result_stack torch.stack([x for _ in range(repeat_times)], dim0) # [3, 100, 200] # 方法3: cat (如果新維度不是第一維stack更方便) # 例如想在最后一維后面拼接 result_cat torch.cat([x.unsqueeze(-1) for _ in range(repeat_times)], dim-1) # [100, 200, 3] result_cat_alt x.unsqueeze(-1).repeat(1, 1, repeat_times) # 效果相同 print(fRepeat correct shape: {result_repeat_correct.shape}) print(fStack shape: {result_stack.shape}) print(fCat shape: {result_cat.shape}) print(fAre stack and repeat equal? {torch.allclose(result_stack, result_repeat_correct)})stack會(huì)創(chuàng)建一個(gè)新的維度而cat是在已有的維度上拼接。選擇哪一個(gè)取決于你想要的輸出形狀。6.3 避免在循環(huán)中調(diào)用repeat如果你需要在循環(huán)中多次對(duì)一個(gè)張量進(jìn)行重復(fù)操作最好在循環(huán)外先計(jì)算好最終形狀一次性分配內(nèi)存并填充而不是在每次迭代中重復(fù)調(diào)用repeat()后者會(huì)帶來(lái)大量的內(nèi)存分配和拷貝開(kāi)銷(xiāo)。# 不推薦的做法 results [] base_tensor torch.randn(10, 20) for i in range(1000): # 每次循環(huán)都創(chuàng)建一個(gè)新的張量 repeated base_tensor.repeat(2, 1) # 形狀 [20, 20] # ... 一些處理 ... results.append(repeated) final_result torch.stack(results, dim0) # 最終形狀 [1000, 20, 20] # 推薦的做法預(yù)分配內(nèi)存 batch_size 1000 target_shape (batch_size, 20, 20) preallocated torch.empty(target_shape, dtypebase_tensor.dtype, devicebase_tensor.device) # 使用 broadcasting 或 view expand 進(jìn)行填充 # 例如如果每個(gè)批次都是 base_tensor 的重復(fù) base_expanded base_tensor.unsqueeze(0).expand(batch_size, -1, -1) # [1000, 10, 20]? 不對(duì)我們需要[1000,20,20] # 更合適的做法是如果邏輯允許直接生成最終數(shù)據(jù)避免中間的 repeat 操作。對(duì)于這種模式如果可能應(yīng)重新思考數(shù)據(jù)流看是否能使用向量化操作避免循環(huán)。如果必須循環(huán)至少確保repeat操作不在最內(nèi)層循環(huán)。理解torch.repeat()的關(guān)鍵在于把握其“維度優(yōu)先”的擴(kuò)展邏輯和物理復(fù)制的本質(zhì)。在大多數(shù)需要數(shù)據(jù)副本的場(chǎng)景下它是得力的工具。但在追求極致性能和大數(shù)據(jù)處理的場(chǎng)合務(wù)必審視是否有更節(jié)省內(nèi)存的替代方案例如expand()或重構(gòu)數(shù)據(jù)處理流程。通過(guò)本文的詳細(xì)拆解和對(duì)比希望你能在下次使用repeat()時(shí)更加得心應(yīng)手精準(zhǔn)地控制張量的形狀與內(nèi)存。