Pytorch-在softmax层之后选择最佳概率 [英] Pytorch - Pick best probability after softmax layer

问题描述

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

I have a logistic regression model using Pytorch 0.4.0, where my input is high-dimensional and my output must be a scalar - 0, 1 or 2.

我正在使用结合softmax层的线性层来返回n x 3张量，其中每一列代表输入落入三类之一(0，1或).

I'm using a linear layer combined with a softmax layer to return a n x 3 tensor, where each column represents the probability of the input falling in one of the three classes (0, 1 or 2).

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

However, I must return a n x 1 tensor, so I need to somehow pick the highest probability for each input and create a tensor indicating which class had the highest probability. How can I achieve this using Pytorch?

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

To illustrate, my Softmax outputs this:

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

我必须返回这个:

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

推荐答案

torch.argmax() 可能就是您想要的:

torch.argmax() is probably what you want:

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层之后选择最佳概率的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！