我有一个使用Pytorch0.4.0的逻辑回归模型,其中我的输入是高维的,并且我的输出必须是标量- 0
、1
或2
。
我使用线性层和softmax层相结合来返回n x 3
张量,其中每一列表示输入落入三个类别(0
、1
或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]]
发布于 2018-06-10 01:01:52
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]])
https://stackoverflow.com/questions/50776548
复制相似问题