argmax函数的使用(pytorch版)
如果一个数组中最大值个数不唯一。那么返回第一个最大值的索引。
·
import torch
input = torch.tensor([[6, 0],
[2, 9]])
print(input.argmax())
# tensor(3) （把[[6,0],[2,9]]看成[6,0,2,9]。从0索引，第3个索引的元素最大）
print(input.argmax(0))
# tensor([0, 1])
# 沿列看，第一列第0个位置（元素'6'）最大，
# 沿列看，第二列第1个位置（元素'9'）最大。依旧是从0开始索引）
print(input.argmax(1))
# tensor([0, 1])
# 沿行看，第一行第0个位置（元素'6'）最大，
# 沿行看，第二行第1个位置（元素'9'）最大。依旧是从0开始索引）
如果一个数组中最大值个数不唯一
那么返回第一个最大值的索引
input2 = torch.tensor([0,6,9,2,1,9])
print(input2.argmax())
# tensor(2) (返回第一个最大值的索引)
魔乐社区（Modelers.cn) 是一个中立、公益的人工智能社区，提供人工智能工具、模型、数据的托管、展示与应用协同服务，为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作，由全产业链共同建设、共同运营、共同享有，推动国产AI生态繁荣发展。
更多推荐
0
0- 0
已为社区贡献12条内容
所有评论(0)