pytorch中torch.stack()函數(shù)用法解讀

更新時間：2024年04月24日 08:54:32 作者：RealWeakCoder

這篇文章主要介紹了pytorch中torch.stack()函數(shù)用法,具有很好的參考價值,希望對大家有所幫助,如有錯誤或未考慮完全的地方,望不吝賜教

torch.stack()函數(shù)用法

一、基本功能

pytroch官方文檔對于這個函數(shù)的描述很簡略。

只有一句話：

在維度上連接（concatenate）若干個張量。(這些張量形狀相同）。

經(jīng)過代碼總結(jié)歸納，可以得到stack（tensors,dim=0,out=None）函數(shù)的功能：

將若干個張量在dim維度上連接,生成一個擴(kuò)維的張量，比如說原來你有若干個2維張量，連接可以得到一個3維的張量。

設(shè)待連接張量維度為n，dim取值范圍為-n-1~n，這里得提一下為負(fù)的意義：-i為倒數(shù)第i個維度。

舉個例子：

對于2維的待連接張量，-1維即3維，-2維即2維。

上代碼：

a=torch.tensor([[1,2,3],[4,5,6]])
b=torch.tensor([[10,20,30],[40,50,60]])
c=torch.tensor([[100,200,300],[400,500,600]])
print(torch.stack([a,b,c],dim=0))
print(torch.stack([a,b,c],dim=1))
print(torch.stack([a,b,c],dim=2))
print(torch.stack([a,b,c],dim=0).size())
print(torch.stack([a,b,c],dim=1).size())
print(torch.stack([a,b,c],dim=2).size())
#輸出結(jié)果為：
tensor([[[  1,   2,   3],
         [  4,   5,   6]],

        [[ 10,  20,  30],
         [ 40,  50,  60]],

        [[100, 200, 300],
         [400, 500, 600]]])
tensor([[[  1,   2,   3],
         [ 10,  20,  30],
         [100, 200, 300]],

        [[  4,   5,   6],
         [ 40,  50,  60],
         [400, 500, 600]]])
tensor([[[  1,  10, 100],
         [  2,  20, 200],
         [  3,  30, 300]],

        [[  4,  40, 400],
         [  5,  50, 500],
         [  6,  60, 600]]])
torch.Size([3, 2, 3])
torch.Size([2, 3, 3])
torch.Size([2, 3, 3])

二、規(guī)律分析

通過代碼運(yùn)行結(jié)果，我們不難發(fā)現(xiàn)，stack（tensors,dim=0,out=None）函數(shù)的運(yùn)行機(jī)制可以等價為：

dim=0時，將tensor在一維上連接，簡單來說就是，就是將tensor1，tensor2…tensor n,連接為【tensor1，tensor2… tensor n】（就是在這里產(chǎn)生了擴(kuò)維）
dim=1時，將每個tensor的第i行按行連接組成一個新的2維tensor，再將這些新tensor按照dim=0的方式連接。
dim=2時，將每個tensor的第i行轉(zhuǎn)置后按列連接組成一個新的2維tensor，再將這些新tesnor按照dim=0的方式連接

可以得到一個結(jié)論:n維（n>=2）待連接張量按dim=x的方式連接等價于：

若x=0,參照上面的規(guī)律進(jìn)行連接
若x>0,對每個張量的第一個維度下的張量對應(yīng)地按照dim=x-1的方式進(jìn)行連接得到若干個新張量,這些新張量按照dim=0的方式進(jìn)行連接。
很明顯，該規(guī)律具有遞歸的特性，x=0,1,2的基礎(chǔ)情況已經(jīng)給出。

注：以上規(guī)律是在未看函數(shù)實現(xiàn)源碼基礎(chǔ)上未加證明的猜測。