PyTorch 深度学习笔记（十一）：Softmax 激活函数在分类任务中的实战

1. Softmax 函数原理

Softmax 是将实数向量映射为概率分布的核心函数，定义为： $$ \sigma(\mathbf{z})i = \frac{e^{z_i}}{\sum{j=1}^{K} e^{z_j}} $$ 其中：

• $\mathbf{z}$ 为模型输出的原始分数向量
• $K$ 表示分类类别数
• 输出满足 $\sum_{i=1}^{K} \sigma(\mathbf{z})_i = 1$

2. PyTorch 实现要点

2.1 关键模块

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

# 网络最后一层无需激活函数
class Classifier(nn.Module):
    def __init__(self, input_dim, num_classes):
        super().__init__()
        self.fc = nn.Linear(input_dim, num_classes)  # 直接输出logits
        
    def forward(self, x):
        return self.fc(x)  # 原始输出值

2.2 Softmax 应用方式

# 方法1：显式调用（验证时使用）
probs = F.softmax(logits, dim=1)

# 方法2：隐式结合损失函数（训练推荐）
loss_fn = nn.CrossEntropyLoss()  # 内置Softmax计算
loss = loss_fn(logits, labels)

3. 分类任务实战（以MNIST为例）

# 数据准备
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_loader = DataLoader(
    datasets.MNIST('./data', train=True, download=True, transform=transform),
    batch_size=64, shuffle=True)

# 模型配置
model = Classifier(input_dim=784, num_classes=10)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练循环
for epoch in range(5):
    for images, labels in train_loader:
        images = images.view(-1, 28*28)
        logits = model(images)
        
        loss = F.cross_entropy(logits, labels)  # Softmax+负对数似然
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

4. 预测与评估

# 预测概率分布
with torch.no_grad():
    logits = model(test_images)
    probs = F.softmax(logits, dim=1)
    preds = torch.argmax(probs, dim=1)

# 计算准确率
accuracy = (preds == test_labels).float().mean()
print(f'测试准确率: {accuracy.item():.4f}')

5. 核心注意事项

1. 数值稳定性
   PyTorch 的 cross_entropy 使用 $log(softmax(\cdot))$ 的优化实现，避免 $e^{z_i}$ 溢出问题

2. 维度指定
   Softmax 需明确 dim 参数：

   • dim=0：按列计算
   • dim=1：按行计算（批处理数据默认）

3. 与 Sigmoid 区别

   特性	Softmax	Sigmoid
   适用场景	互斥多分类 ($K>2$)	二分类/多标签
   输出关系	概率和为1	独立概率

4. 温度系数 $\tau$
   可通过调节温度参数控制分布平滑度：

   scaled_probs = F.softmax(logits / tau, dim=1)
   

   当 $\tau \to \infty$ 时输出趋近均匀分布，$\tau \to 0^+$ 时趋近one-hot分布

6. 可视化决策边界

# 生成测试网格
xx, yy = np.meshgrid(np.linspace(-3, 3, 100), np.linspace(-3, 3, 100))
grid = torch.tensor(np.c_[xx.ravel(), yy.ravel()]).float()

# 预测概率
probs = model(grid).softmax(dim=1)
class_idx = torch.argmax(probs, dim=1).numpy().reshape(xx.shape)

# 绘制等高线
plt.contourf(xx, yy, class_idx, alpha=0.5)
plt.scatter(X[:,0], X[:,1], c=y, edgecolors='k')
plt.title('Softmax决策边界')

实战建议：在图像分类等任务中，优先使用 nn.CrossEntropyLoss 而非显式Softmax层，可提升计算效率并避免数值误差累积。