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