argmax函數(shù)：torch.argmax(input, dim=None, keepdim=False) 返回指定維度最大值的序號(hào)，dim給定的定義是：the demention to reduce，就是把dim這個(gè)維度，變成這個(gè)維度的最大值的index。

1）dim表示不同維度。特別的在dim=0表示二維矩陣中的列，dim=1在二維矩陣中的行。廣泛的來(lái)說(shuō)，我們不管一個(gè)矩陣是幾維的，比如一個(gè)矩陣維度如下：(d0,d1,…,dn?1) ，那么dim=0就表示對(duì)應(yīng)到d0 也就是第一個(gè)維度，dim=1表示對(duì)應(yīng)到也就是第二個(gè)維度，以此類推。

2）知道dim的值是什么意思還不行，還要知道函數(shù)中這個(gè)dim給出來(lái)會(huì)發(fā)生什么。

例子一：二維數(shù)組

import torchx = torch.randn(2, 4) print(x) ''' tensor([[ 1.2864, -0.5955, 1.5042, 0.5398],[-1.2048, 0.5106, -2.0288, 1.4782]]) '''# y0表示矩陣dim=0維度上（每一列）張量最大值的索引 y0 = torch.argmax(x, dim=0) print(y0) ''' tensor([0, 1, 0, 1]) '''# y1表示矩陣dim=1維度上（每一行）張量最大值的索引 y1 = torch.argmax(x, dim=1) print(y1) ''' tensor([2, 3]) '''

例子二：三維數(shù)組

x = torch.randn(2, 4, 5) print(x) ''' tensor([[[-1.2204, -0.6428, -0.2278, 0.5589, 1.1589],[ 0.4235, 1.9663, 0.5055, -1.3472, 1.3523],[ 1.4220, 0.7886, -1.0821, 0.6268, -0.9465],[-0.3950, 1.3275, 0.3369, 1.0224, -0.9944]],[[ 0.6024, -0.2604, -0.8631, 0.8113, -0.3140],[ 0.3487, -0.1941, -0.3955, -0.1719, -1.3734],[ 0.2467, -0.4268, -1.3428, 0.7346, 1.0932],[-0.5799, 0.0976, -1.9403, -0.2643, 0.7657]]]) '''# dim=0，將第一個(gè)維度消除，也就是將兩個(gè)[4*5]矩陣只保留一個(gè)，因此要在上下兩個(gè)[3*4]的矩陣分別在對(duì)應(yīng)位置上比較 y0 = torch.argmax(x, dim=0) print(y0) ''' tensor([[1, 1, 0, 1, 0],[0, 0, 0, 1, 0],[0, 0, 0, 1, 1],[0, 0, 0, 0, 1]]) '''# dim=1，將第二個(gè)維度消除，也就是將四個(gè)[2*5]矩陣只保留一個(gè) y1 = torch.argmax(x, dim=1) print(y1) ''' tensor([[2, 1, 1, 3, 1],[0, 3, 1, 0, 2]]) '''y2 = torch.argmax(x, dim=2) print(y2) ''' tensor([[4, 1, 0, 1],[3, 0, 4, 4]]) '''

總結(jié)

以上是生活随笔為你收集整理的torch.argmax()函数的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

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