pytorch中max和argmax

max和argmax名字上看起来很相似但是却大有不同

max_num,index = torch.max(tensor,dim=?)会在指定的维度上返回最大的数以及其索引如下

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)
max_num,index = torch.max(b,dim=2)
print(max_num)
print(index)

输出：

tensor([[[3, 2],
         [1, 4]],

        [[5, 6],
         [7, 8]]])
**************************************************
tensor([[3, 4],
        [6, 8]])
**************************************************
tensor([[0, 1],
        [1, 1]])

 index = torch.argmax(tensor,dim=?)是在指定维度上输出最大值的索引，代码如下

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)
print('*'*50)
index = torch.argmax(b,dim=2)

输出：

输出：
tensor([[[3, 2],
         [1, 4]],

        [[5, 6],
         [7, 8]]])
*****************************