PyTorch 中對(duì) tensor 的很多操作如?sum、softmax?等都可以設(shè)置?dim?參數(shù)用來(lái)指定操作在哪一維進(jìn)行。PyTorch 中的 dim 類似于 numpy 中的 axis，這篇文章來(lái)總結(jié)一下 PyTorch 中的 dim 操作。首先看一下這個(gè)圖，圖中給出了維度標(biāo)號(hào)，注意區(qū)分正負(fù)，從左往右數(shù)，括號(hào)代表的維度分別是 0 和 1 和 2，從右往做為 -3 和 -2 和 -1。待會(huì)兒會(huì)用到。

圖1

括號(hào)之間是嵌套關(guān)系，代表了不同的維度。從左往右數(shù)，兩個(gè)括號(hào)代表的維度分別是 0 和 1 ，在第 0 維遍歷得到向量，在第 1 維遍歷得到標(biāo)量.

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

則 3 個(gè)括號(hào)代表的維度從左往右分別為 0, 1, 2，在第 0 維遍歷得到矩陣，在第 1 維遍歷得到向量，在第 2 維遍歷得到標(biāo)量。

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

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

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

求 a 在第 0 維的和，因?yàn)榈?0 維代表最外邊的括號(hào)，括號(hào)中的元素為向量[1, 2]，[3, 4]，第 0 維的和就是第 0 維中的元素相加，也就是兩個(gè)向量[1, 2]，[3, 4]相加，所以結(jié)果為[4,6]

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

輸出

tensor([4, 6])

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

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

輸出

tensor([[4, 6]])

同理的現(xiàn)在對(duì)dim=1進(jìn)行操作

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]])

現(xiàn)在對(duì)三維張量進(jìn)行操作

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 維為最外層括號(hào)，最外層括號(hào)中的元素為矩陣[[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 維求和，就是對(duì)標(biāo)量 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]]])

現(xiàn)在再來(lái)看看其他dim有關(guān)的api

1.torch.max

在二維中

dim = 0

這個(gè)時(shí)候取的是矩陣的最大值以及下標(biāo)即[[1,2],[3,4]]中的最大值和下標(biāo)，那么應(yīng)該是[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

這個(gè)時(shí)候取的是標(biāo)量是在[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

則是對(duì)比這兩個(gè)矩陣返回同位置最大的值

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]這四個(gè)中返回最大值

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

對(duì)標(biāo)量做比較

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]]))

那么同樣的現(xiàn)在來(lái)思考下dim=-3，-2，-1的情況，圖1中已經(jīng)給出來(lái)了正負(fù)數(shù)維度對(duì)應(yīng)的關(guān)系了，現(xiàn)在我們就取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]]))

發(fā)現(xiàn)了嗎dim=-1和dim=2輸出的是相同的，剩余的可自行驗(yàn)證。

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。