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

PyTorch中求最大值索引的操作演示

原创于 2022-03-21 11:58:43 发布 · 1.1k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#python

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

部署运行你感兴趣的模型镜像

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


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

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch

Cuda

PyTorch 是一个开源的 Python 机器学习库，基于 Torch 库，底层由 C++ 实现，应用于人工智能领域，如计算机视觉和自然语言处理