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

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！