pytorch里的argmax和argmin的理解

返回输入张量中所有元素的最大值/最小值的索引。

torch.argmin(dim=None, keepdim=False) → LongTensor
torch.argmax(input, dim, keepdim=False) → LongTensor

dim：用于指定在哪个维度进行操作
keepdim：若为True,则输出的索引的维数和输入的张量的维度保持一致，一般默认为False

实验如下

a = torch.randn(4, 4)
print(a)
print(a.flatten())
print(torch.argmax(a))  # 不指定dim，则返回张量展开为一维最大值的索引
"""
tensor([[ 1.0284, -0.1810,  2.2535,  1.0803],
        [ 0.0180,  0.8234, -1.8942, -0.0986],
        [-0.7464,  0.5619,  1.1605,  0.3683],
        [-0.7363, -0.2394, -0.6425, -0.5612]])
tensor([ 1.0284, -0.1810,  2.2535,  1.0803,  0.0180,  0.8234, -1.8942, -0.0986,
        -0.7464,  0.5619,  1.1605,  0.3683, -0.7363, -0.2394, -0.6425, -0.5612])
tensor(2)
"""
a = torch.randn(4, 4)
print(a)
print(torch.argmax(a), dim=1)
"""
tensor([[-0.3381, -0.4999, -1.1872, -0.6996],
        [-1.1637, -1.9946, -1.5603,  3.1711],
        [ 0.1105, -1.4377,  0.0382, -0.3079],
        [ 2.2044,  0.9992, -1.6193, -0.2622]])
tensor([0, 3, 0, 0])
"""