PyTorch 中對 tensor 的很多操作如?sum、softmax?等都可以設置?dim?參數用來指定操作在哪一維進行。PyTorch 中的 dim 類似于 numpy 中的 axis，這篇文章來總結一下 PyTorch 中的 dim 操作。首先看一下這個圖，圖中給出了維度標號，注意區分正負，從左往右數，括號代表的維度分別是 0 和 1 和 2，從右往做為 -3 和 -2 和 -1。待會兒會用到。

圖1

括號之間是嵌套關系，代表了不同的維度。從左往右數，兩個括號代表的維度分別是 0 和 1 ，在第 0 維遍歷得到向量，在第 1 維遍歷得到標量.

a = torch.tensor([[1,2],[3,4]])

則 3 個括號代表的維度從左往右分別為 0, 1, 2，在第 0 維遍歷得到矩陣，在第 1 維遍歷得到向量，在第 2 維遍歷得到標量。

b = torch.tensor([[[3, 2], [1, 4]],[[5, 6], [7, 8]]])#張量

在某一維度求和（或者進行其他操作）就是對該維度中的元素進行求和。對于矩陣 a

a = torch.tensor([[1,2],[3,4]])

求 a 在第 0 維的和，因為第 0 維代表最外邊的括號，括號中的元素為向量[1, 2]，[3, 4]，第 0 維的和就是第 0 維中的元素相加，也就是兩個向量[1, 2]，[3, 4]相加，所以結果為[4,6]

s = torch.sum(a, dim=0) print(s)

輸出

tensor([4, 6])

可以看到，a 是 2 維矩陣，而相加的結果為 1 維向量，可以使用參數keepdim=True來保證維度數目不變。

s = torch.sum(a, dim=0, keepdim=True) print(s)

輸出

tensor([[4, 6]])

同理的現在對dim=1進行操作

a = torch.tensor([[1,2],[3,4]]) s = torch.sum(a,dim=1,keepdim=False) print(s) #輸出 tensor([3, 7])

keepdim = True

a = torch.tensor([[1,2],[3,4]]) s = torch.sum(a,dim=1,keepdim=True) print(s) # 輸出 tensor([[3],[7]])

現在對三維張量進行操作

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) print(b) # 輸出 tensor([[[3, 2],[1, 4]],[[5, 6],[7, 8]]])

將 b 在第 0 維相加，第 0 維為最外層括號，最外層括號中的元素為矩陣[[3, 2], [1, 4]]和[[5, 6], [7, 8]]。在第 0 維求和，就是將第 0 維中的元素（矩陣）相加

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) #print(b) s = torch.sum(b,dim=0) print(s) # 輸出 tensor([[ 8, 8],[ 8, 12]]) keepdim = True #輸出tensor([[[ 8, 8],[ 8, 12]]])

求 b 在第 1 維的和，就是將 b 第 1 維中的元素[3, 2]和[1, 4],?[5, 6]和?[7, 8]相加，所以

[3,2]+[1,4]=[4,6],[5,6]+[7,8]=[12,14]

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) #print(b) s = torch.sum(b,dim=1) print(s) #輸出 tensor([[ 4, 6],[12, 14]]) keepdim = True #輸出 tensor([[[ 4, 6]],[[12, 14]]])

則在 b 的第 2 維求和，就是對標量 3 和 2, 1 和 4, 5 和 6 , 7 和 8 求和

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) #print(b) s = torch.sum(b,dim=2) print(s) #輸出 tensor([[ 5, 5],[11, 15]]) keepdim = True #輸出 tensor([[[ 5],[ 5]],[[11],[15]]])

現在再來看看其他dim有關的api

1.torch.max

在二維中

dim = 0

這個時候取的是矩陣的最大值以及下標即[[1,2],[3,4]]中的最大值和下標，那么應該是[3,4]

a = torch.tensor([[1,2],[3,4]]) print(a) print(torch.max(a,dim=0)) #輸出 values=tensor([3, 4]), indices=tensor([1, 1]))

dim = 1

這個時候取的是標量是在[1,2]和[3,4]中找到最大值

a = torch.tensor([[1,2],[3,4]]) print(a) print(torch.max(a,dim=1)) # 輸出 values=tensor([2, 4]), indices=tensor([1, 1]))

在三維中

dim = 0

則是對比這兩個矩陣返回同位置最大的值

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) print(b) print(torch.max(b,dim=0)) # 輸出 tensor([[[3, 2],[1, 4]],[[5, 6],[7, 8]]]) torch.return_types.max( values=tensor([[5, 6],[7, 8]]), indices=tensor([[1, 1],[1, 1]]))

dim = 1

[3, 2], [1, 4]，[5, 6], [7, 8]這四個中返回最大值

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) print(b) print(torch.max(b,dim=1)) # 輸出 tensor([[[3, 2],[1, 4]],[[5, 6],[7, 8]]]) torch.return_types.max( values=tensor([[3, 4],[7, 8]]), indices=tensor([[0, 1],[1, 1]]))

dim = 2

對標量做比較

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) print(b) print(torch.max(b,dim=2)) # 輸出 tensor([[[3, 2],[1, 4]],[[5, 6],[7, 8]]]) torch.return_types.max( values=tensor([[3, 4],[6, 8]]), indices=tensor([[0, 1],[1, 1]]))

那么同樣的現在來思考下dim=-3，-2，-1的情況，圖1中已經給出來了正負數維度對應的關系了，現在我們就取dim=-1看看

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]]) c = torch.max(b,dim=-1) print(c)

輸出：

torch.return_types.max( values=tensor([[3, 4],[6, 8]]), indices=tensor([[0, 1],[1, 1]]))

發現了嗎dim=-1和dim=2輸出的是相同的，剩余的可自行驗證。

總結

以上是生活随笔為你收集整理的pytorch中维度dim的理解的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。