torch.stack()

如圖有a, b, c 三個 3×3的Tensor,

如果我想把這三個tensor的最後一個維度的元素相疊加,形成一個新的tensor

輸入 d=torch.stack( (a,b,c) ,dim = 2)

就會發現他們在第二個維度像疊加,具體來說變化如下圖所示

d[0][0]的位置是由a[0][0]  的 [1]   , b[0][0] 的 [10],  c[0][0] 的 [100] ,這3個 元素疊加構成的一個size為3的新的元素[1,10,100],所以d的維度就是3 * 3  *3。這裡注意dim參數!!!!!!!雖然他們是在第二個維度相疊加,(pytorch中從0開始計數,所以也就是第一個維度是dim=0,第二個維度是dim=1),但是我在stack函數中寫的dim卻是 dim=2,也即第三個維度,這是說我們最終的生成結果位於第三個維度,而不是說,原來的元素是第三個維度。所以還有另外一種寫法,寫dim=-1,不管你原來是啥,我就指定最後一個維度。

得到的結果和我一開始寫dim=2是一樣。

再補充一個例子讓各位感受一下這個dim參數的意義把,它指定的是做完stack後的元素結果所在的那。

dim分別設置為0 ,1,2  注意觀察結果

c, dim = 0時, c = [ a, b]

d, dim =1 時, d = [ [a[0] , b[0] ] , [a[1], b[1] ] ]

e, dim = 2 時, e = [     [   [ a[0][0], b[0][0] ]  , [ a[0][1], b[0][1] ]  ,  [ a[0][2],b[0][2] ] ] ,                                    

 [   [ a[1][0], b[1][0] ]  , [ a[1][1], b[0][1] ]  ,  [ a[1][2],b[1][2] ] ]      ]