我们构造一个简单的神经网络，通常情况下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]])
