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]]])
**************************************************
tensor([[0, 1],
        [1, 1]])

同样的我们可以看看min和argmin

torch.min

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

输出：

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

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

torch.argmin

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

输出：

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

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