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

Stack Overflow用户

提问于 2018-06-10 00:55:10

回答 1查看 7.1K关注 0票数 5

我有一个使用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]]

页面原文内容由Stack Overflow提供。腾讯云小微IT领域专用引擎提供翻译支持

原文链接：

https://stackoverflow.com/questions/50776548

复制

相似问题

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