想要神经网络输出分类的概率值？应该这样写代码

我们构造一个简单的神经网络，通常情况下n_output是分类数量，例如二分类任务那n_output=2、六分类任务那么n_output=6

class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        self.inLayer = torch.nn.Linear(n_feature, n_hidden) # 输入层
        self.hiddenLayer = torch.nn.Linear(n_hidden, n_hidden) # 隐藏层
        self.outLayer = torch.nn.Linear(n_hidden, n_output) # 输出层

    # 前向计算函数，定义完成后会隐式地自动生成反向传播函数
    def forward(self, x):
        x = F.relu(self.hiddenLayer(self.inLayer(x)))
        x = self.outLayer(x)
        return x

假设要进行一个二分类任务，这时的输出x的值是类别0和类别1的预测值，我们可以使用softmax函数将其转换为概率值：

class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        self.inLayer = torch.nn.Linear(n_feature, n_hidden) # 输入层
        self.hiddenLayer = torch.nn.Linear(n_hidden, n_hidden) # 隐藏层
        self.outLayer = torch.nn.Linear(n_hidden, n_output) # 输出层

    # 前向计算函数，定义完成后会隐式地自动生成反向传播函数
    def forward(self, x):
        x = F.relu(self.hiddenLayer(self.inLayer(x)))
        x = self.outLayer(x)
        prob = F.softmax(x, dim=1) # softmax将预测值转换为概率
        out = torch.unsqueeze(prob.argmax(dim=1), dim=1) # argmax选取概率最大的预测项，unsqueeze扩维
        return out

输出结果例如：

# prob = F.softmax(x, dim=1)的输出结果
# print(prob)
tensor([[0.5059, 0.4941],
        [0.5018, 0.4982],
        [0.4886, 0.5114]], grad_fn=<SoftmaxBackward0>)

# out = torch.unsqueeze(prob.argmax(dim=1), dim=1)的输出结果
# print(out)
tensor([[0],
        [0],
        [1]])