深入理解Softmax函数及其在PyTorch中的实现-使用PyTorch实现Softmax函数

在PyTorch中，可以通过多种方式实现Softmax函数。以下将通过示例演示如何使用torch.nn.functional.softmax和torch.nn.Softmax。

创建输入数据

首先，创建一个示例输入张量：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 创建一个输入张量，形状为 (batch_size, features)
input_tensor = torch.tensor([[2.0, 1.0, 0.1],
                             [1.0, 3.0, 0.2]])
print("输入张量：")
print(input_tensor)

输出：

输入张量：
tensor([[2.0000, 1.0000, 0.1000],
        [1.0000, 3.0000, 0.2000]])

---

方法一：使用torch.nn.functional.softmax

利用PyTorch中torch.nn.functional.softmax函数直接对输入数据应用Softmax。

# 在维度1上（即特征维）应用Softmax
softmax_output = F.softmax(input_tensor, dim=1)
print("\nSoftmax输出：")
print(softmax_output)

输出：

Softmax输出：
tensor([[0.6590, 0.2424, 0.0986],
        [0.1065, 0.8726, 0.0209]])

---

方法二：使用torch.nn.Softmax模块

也可以使用torch.nn中的Softmax模块。

# 创建一个Softmax层实例
softmax = nn.Softmax(dim=1)

# 对输入张量应用Softmax层
softmax_output_module = softmax(input_tensor)
print("\n使用nn.Softmax模块的输出：")
print(softmax_output_module)

输出：

使用nn.Softmax模块的输出：
tensor([[0.6590, 0.2424, 0.0986],
        [0.1065, 0.8726, 0.0209]])

---

在神经网络模型中应用Softmax

构建一个简单的神经网络模型，在最后一层使用Softmax激活函数。

class SimpleNetwork(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SimpleNetwork, self).__init__()
        self.layer1 = nn.Linear(input_size, 5)
        self.layer2 = nn.Linear(5, num_classes)
        # 使用LogSoftmax提高数值稳定性
        self.softmax = nn.LogSoftmax(dim=1)
    
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = self.layer2(x)
        x = self.softmax(x)
        return x

# 定义输入大小和类别数
input_size = 3
num_classes = 3

# 创建模型实例
model = SimpleNetwork(input_size, num_classes)

# 查看模型结构
print("\n模型结构：")
print(model)

输出：

模型结构：
SimpleNetwork(
  (layer1): Linear(in_features=3, out_features=5, bias=True)
  (layer2): Linear(in_features=5, out_features=3, bias=True)
  (softmax): LogSoftmax(dim=1)
)

前向传播：

# 将输入数据转换为浮点型张量
input_data = input_tensor.float()

# 前向传播
output = model(input_data)
print("\n模型输出（对数概率）：")
print(output)

输出：

模型输出（对数概率）：
tensor([[-1.2443, -0.7140, -1.2645],
        [-1.3689, -0.6535, -1.5142]], grad_fn=<LogSoftmaxBackward0>)

转换为概率：

# 取指数，转换为概率
probabilities = torch.exp(output)
print("\n模型输出（概率）：")
print(probabilities)

输出：

模型输出（概率）：
tensor([[0.2882, 0.4898, 0.2220],
        [0.2541, 0.5204, 0.2255]], grad_fn=<ExpBackward0>)

预测类别：

# 获取每个样本概率最大的类别索引
predicted_classes = torch.argmax(probabilities, dim=1)
print("\n预测的类别：")
print(predicted_classes)

输出：

预测的类别：
tensor([1, 1])

---