PyTorch 深度学习笔记（十一）：多分类任务中 Softmax 激活函数的场景适配

在多分类任务中，Softmax 激活函数是神经网络输出层的关键组件。它将原始输出值转换为概率分布，使模型能够预测样本属于每个类别的概率。以下从原理、数学表达和场景适配三个维度展开说明：

1. 核心原理与数学表达

给定神经网络的原始输出向量$z = [z_1, z_2, \ldots, z_K]^T$（$K$为类别数），Softmax 函数定义为： $$ \sigma(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} \quad \text{for} \quad i = 1,2,\ldots,K $$ 该函数满足：

• 输出值域$(0,1)$
• 所有类别概率和为$1$：$\sum_{i=1}^{K} \sigma(z_i) = 1$
• 指数运算放大较大值的影响，突出主导类别

2. 场景适配关键点

3. PyTorch 实现示例

import torch
import torch.nn as nn

class MulticlassClassifier(nn.Module):
    def __init__(self, input_dim, num_classes):
        super().__init__()
        self.fc = nn.Linear(input_dim, num_classes)  # 最后一层线性输出
        self.softmax = nn.Softmax(dim=1)  # 按行计算概率分布
        
    def forward(self, x):
        logits = self.fc(x) 
        return self.softmax(logits)  # 输出概率分布

# 训练流程
model = MulticlassClassifier(input_dim=128, num_classes=10)
criterion = nn.CrossEntropyLoss()  # 内置Softmax计算优化
optimizer = torch.optim.Adam(model.parameters())

# 预测示例
with torch.no_grad():
    probs = model(input_data)  # 获得各类别概率
    pred_class = torch.argmax(probs, dim=1)  # 取最大概率类别

4. 特殊场景适配技巧

• 类别不平衡：在损失函数中引入类别权重

  weights = torch.tensor([0.1, 0.2, 0.7])  # 各类别权重
  criterion = nn.CrossEntropyLoss(weight=weights)
  

• 高维输出：使用dim参数明确指定计算维度

  # 处理三维输出 (batch, seq_len, classes)
  softmax = nn.Softmax(dim=2)
  

• 部署优化：推理时用LogSoftmax替代避免重复计算

  self.log_softmax = nn.LogSoftmax(dim=1)
  

关键提示：当使用CrossEntropyLoss时，无需显式添加Softmax层，因其已包含优化后的数值计算。仅在需要直接获取概率值时显式调用。

通过上述适配方案，Softmax 可有效解决多分类任务中的概率归一化问题，并与PyTorch的损失函数体系无缝集成，是分类任务输出层的标准配置。