Pytorch - 在softmax层之后选择最佳概率

内容来源于 Stack Overflow,并遵循CC BY-SA 3.0许可协议进行翻译与使用

  • 回答 (1)
  • 关注 (0)
  • 查看 (1095)

我一直在使用Pytorch 0.4.0,我的输入为高维和我的输出必须是标量逻辑回归模型- 012

我正在使用一个与softmax图层相结合的线性图层来返回一个n x 3张量,其中每列代表输入落入三个类别(012)之一的概率。

但是,我必须返回一个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]])

扫码关注云+社区

领取腾讯云代金券