pred_classes = torch.max(pred, dim=1)[1]

我竟然有3090

于 2022-03-21 11:58:43 发布

阅读量1k

点赞数 4

CC 4.0 BY-SA版权

文章标签： python

本文链接：https://blog.youkuaiyun.com/STANDEDS/article/details/123631940

这段内容展示了如何在PyTorch中使用`torch.max()`函数找到二维张量沿指定维度的最大值及其对应的索引。示例中，输入是一个3x4的张量，通过`dim=1`参数获取每行的最大值及其索引，最终返回的是最大值和对应的列索引。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

>>> import torch

>>> i=torch.Tensor([[1, 1, 0, 1], 
                [0, 1, 1, 1],
                [0, 1, 0, 1]])

>>> i.size()
torch.Size([3, 4])

>>> c=torch.max(i,dim=1)
>>> c
torch.return_types.max(
values=tensor([1., 1., 1.]),
indices=tensor([0, 1, 1]))

>>> c=torch.max(i,dim=1)[1]
>>> c
tensor([0, 1, 1])


#最终结果就是返回最大值的索引值