关于argmax()

最新推荐文章于 2022-11-06 21:51:12 发布

反正没几根头发

最新推荐文章于 2022-11-06 21:51:12 发布

阅读量375

点赞数

分类专栏： pytorch学习文章标签：深度学习 python pytorch

本文链接：https://blog.csdn.net/weixin_44334635/article/details/126878627

版权

pytorch学习专栏收录该内容

4 篇文章 1 订阅

订阅专栏

input.argmax(1)：横向比较
input.argmax(0)：纵向比较

import torch

input = torch.tensor([[0.6,0.5,0.3],
                      [0.2,0.3,0.1]])

target1 = torch.tensor([0,1])
target2 = torch.tensor([0,0,1])

# 参数为1，横向比较
output1 = input.argmax(1)
print(output1)      # torch.tensor([0,1])
print((output1 == target1).sum())   # 输出预测正确次数：tensor(2)

# 参数为0，纵向比较
output2 = input.argmax(0)
print(output2)      # torch.tensor([0,0,0])
print((output2 == target2).sum())   # 输出预测正确次数：tensor(2)