我有一个使用 Pytorch 0.4.0 的逻辑回归模型,其中我的输入是高维的,我的输出必须是标量 -0
, 1
or 2
.
我使用线性层与 softmax 层相结合来返回n x 3
张量,其中每列表示输入属于三个类别之一的概率 (0
, 1
or 2
).
但是,我必须返回一个n x 1
张量,所以我需要以某种方式为每个输入选择最高概率,并创建一个张量来指示哪个类具有最高概率。我如何使用 Pytorch 实现这一目标?
为了说明这一点,我的 Softmax 输出如下:
[[0.2, 0.1, 0.7],
[0.6, 0.2, 0.2],
[0.1, 0.8, 0.1]]
我必须返回这个:
[[2],
[0],
[1]]
torch.argmax()可能是你想要的:
import torch
x = torch.FloatTensor([[0.2, 0.1, 0.7],
[0.6, 0.2, 0.2],
[0.1, 0.8, 0.1]])
y = torch.argmax(x, dim=1)
print(y.detach())
# tensor([ 2, 0, 1])
# If you want to reshape:
y = y.view(1, -1)
print(y.detach())
# tensor([[ 2, 0, 1]])
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)