在模型最后添加softmax层怎么写代码

在PyTorch中，添加softmax层可以使用nn.Softmax()函数。具体的代码如下所示：

import torch.nn as nn

# 定义模型类
class MyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.softmax = nn.Softmax(dim=1)  # 在dim=1的维度上进行softmax操作

    def forward(self, x):
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        x = self.softmax(x)  # 在输出层添加softmax层
        return x

在上面的代码中，我们在模型的`__init__`方法中定义了一个`nn.Softmax()`层，并且在`forward`方法中将输出结果通过该层进行softmax操作。需要注意的是，`dim=1`表示在第二维度上进行softmax操作，也就是在每个样本的输出结果上进行softmax操作。