torch.max ()与 torch.argmax()的区别

torch.max ()与 torch.argmax()的区别

记录一下：这里之前一直容易搞混，
torch.max ()返回指定dim的最大值和对应的位置索引值，
torch.argmax ()返回指定dim的最大值对应的位置索引值

1. torch.max ()的官方使用教程

2. torch.argmax()的官方使用教程

import torch

a = torch.tensor([[0.03,0.12,0.85],
                  [0.01,0.9,0.09],
                  [0.95,0.01,0.04],
                  [0.09, 0.9, 0.01]])
print(a)
print(a.dtype)

b = torch.max(a, 1)
print('b:', b)
print("b第0维度：", b[0])
print("b第1维度：", b[1])

c = torch.argmax(a, 0)
print('c第0维度:', c)
c = torch.argmax(a, 1)
print('c第1维度:', c)

运行结果如下：

tensor([[0.0300, 0.1200, 0.8500],
        [0.0100, 0.9000, 0.0900],
        [0.9500, 0.0100, 0.0400],
        [0.0900, 0.9000, 0.0100]])
torch.float32
b: torch.return_types.max(
values=tensor([0.8500, 0.9000, 0.9500, 0.9000]),
indices=tensor([2, 1, 0, 1]))
b第0维度： tensor([0.8500, 0.9000, 0.9500, 0.9000])
b第1维度： tensor([2, 1, 0, 1])
c第0维度: tensor([2, 1, 0])
c第1维度: tensor([2, 1, 0, 1])