Pytorch

Question

我有一个使用 Pytorch 0.4.0 的逻辑回归模型，其中我的输入是高维的，我的输出必须是标量 - 0、1 或 2。

我将线性层与 softmax 层结合使用以 return 一个 n x 3 张量，其中每一列代表输入落入三个 class 之一的概率]es（0、1 或 2）。

但是，我必须 return 一个 n x 1 张量，所以我需要以某种方式为每个输入选择最高概率并创建一个张量来指示哪个 class 具有最高概率。我如何使用 Pytorch 实现此目的？

为了说明，我的 Softmax 输出如下：

[[0.2, 0.1, 0.7],
 [0.6, 0.2, 0.2],
 [0.1, 0.8, 0.1]]

而且我必须return这个：

[[2],
 [0],
 [1]]

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

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

Pytorch - Pick best probability after softmax layer

python

numpy

softmax