PyTorch单隐藏层MLP无法训练问题排查求助

问题背景 #

手动实现梯度下降（不使用PyTorch内置优化器）训练单隐藏层MLP，针对MNIST数据集训练时，损失值始终维持在2.3左右，无法有效下降。相关代码如下：

import torch
import numpy as np
import torchvision.datasets as datasets
from torchvision.transforms import ToTensor

# download train and test set
train = datasets.MNIST(root='.data', train=True, download=True, transform=ToTensor())
test = datasets.MNIST(root='.data', train=False, download=True, transform=ToTensor())

# calculate mean and std
mean = train.data.double().mean()
std = train.data.double().std()
print(f"Mean: {mean}, std: {std}")

# standardize data
train.data = (train.data - mean) / std
test.data = (test.data - mean) / std

from torch.utils.data import DataLoader

train_loader1 = DataLoader(train, shuffle=True, batch_size=1)
train_loader2 = DataLoader(train, shuffle=True, batch_size=256)

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten() # flatten the 28x28 images to 784 features vectors
        self.in_to_hidden = nn.Linear(784, 784) # input to 784 neurons hidden layer
        self.hidden_to_out = nn.Linear(784, 10) # 784 to 10 neuron output layer
        self.softmax = nn.Softmax(dim=1) # make the output into a probability distribution
        
    def forward(self, x):
        activation_fn = nn.Tanh()
        x = self.flatten(x)
        x = activation_fn(self.in_to_hidden(x))
        x = self.softmax(self.hidden_to_out(x))
        return x
    
    def train_single_epoch(self, data_loader, loss_fn, lr, device):
        for x, y in data_loader:
            x, y = x.to(device), y.to(device)

            # calculate loss
            pred = self(x)
            loss = loss_fn(pred, y)

            # backpropagate error and update weights
            loss.backward()
            with torch.no_grad(): # the gradient mustn't be calculated for the weight updates
                for weights in self.parameters():
                    weights -= lr * weights.grad
                    weights.grad.zero_()

        print(f"loss: {loss.item()}")

net = SimpleNN().to(device)
for i in range(2):
    net.train_single_epoch(train_loader2, nn.CrossEntropyLoss(), 0.001, device)

关键排查点及修复方案 #

1. CrossEntropyLoss与Softmax的冲突 #

CrossEntropyLoss内部已经整合了LogSoftmax和NLLLoss的计算逻辑，额外添加输出层的Softmax会导致损失计算出现数值稳定性问题，进而引发梯度消失。
修复：移除self.softmax，让hidden_to_out输出的原始logits直接传入损失函数：

def forward(self, x):
    activation_fn = nn.Tanh()
    x = self.flatten(x)
    x = activation_fn(self.in_to_hidden(x))
    x = self.hidden_to_out(x)  # 直接输出logits
    return x

2. 数据标准化未生效 #

直接修改train.data和test.data的方式不会作用到DataLoader取出的数据上——ToTensor()已将数据转换为0-1的张量，而train.data是原始uint8格式数据，两者存储独立。
修复：用transforms.Normalize将标准化集成到预处理流程：

from torchvision import transforms

# 计算0-1范围内的均值和标准差
mean = train.data.double().mean() / 255
std = train.data.double().std() / 255

transform = transforms.Compose([
    ToTensor(),
    transforms.Normalize(mean=[mean], std=[std])
])

# 重新创建数据集
train = datasets.MNIST(root='.data', train=True, download=True, transform=transform)
test = datasets.MNIST(root='.data', train=False, download=True, transform=transform)

3. 梯度清零时机错误 #

在更新每个参数后立刻调用weights.grad.zero_()，会导致后续参数更新时使用的梯度已被清零，参数更新不完整。
修复：在每个batch的前向传播前统一清零所有梯度：

def train_single_epoch(self, data_loader, loss_fn, lr, device):
    for x, y in data_loader:
        x, y = x.to(device), y.to(device)
        
        # 先清零所有梯度
        self.zero_grad()
        
        # 前向传播计算损失
        pred = self(x)
        loss = loss_fn(pred, y)
        
        # 反向传播
        loss.backward()
        
        # 更新参数
        with torch.no_grad():
            for param in self.parameters():
                param -= lr * param.grad

4. 学习率与激活函数适配问题 #

Tanh激活函数容易出现饱和导致梯度消失，当前设置的学习率0.001偏小，不足以驱动参数更新。
修复：将学习率调整到0.01，或者替换为ReLU激活函数：

# 替换激活函数示例
def forward(self, x):
    activation_fn = nn.ReLU()  # 改用ReLU
    x = self.flatten(x)
    x = activation_fn(self.in_to_hidden(x))
    x = self.hidden_to_out(x)
    return x

5. 参数初始化优化 #

默认的Linear层初始化可能不适合Tanh激活，容易导致初始梯度过小。
修复：使用Xavier初始化适配Tanh：

def __init__(self):
    super().__init__()
    self.flatten = nn.Flatten()
    self.in_to_hidden = nn.Linear(784, 784)
    self.hidden_to_out = nn.Linear(784, 10)
    
    # Xavier初始化
    nn.init.xavier_uniform_(self.in_to_hidden.weight)
    nn.init.zeros_(self.in_to_hidden.bias)
    nn.init.xavier_uniform_(self.hidden_to_out.weight)
    nn.init.zeros_(self.hidden_to_out.bias)

内容的提问来源于stack exchange，提问作者A_Pumpkin