手写数字识别神经网络Cross Entropy Loss不收敛问题求助

我是神经网络新手，目前正在搭建一个手写数字识别模型，但训练时损失值始终在2.2-2.3左右波动，无法收敛。我尝试调整学习率但没有效果，相关代码如下：

模型代码 #

# TODO: Define function to create our own neural network

# Parameters
input_size = 784    # Hint: image size is 28x28, and we want to flatten the image 
num_classes = 10   # Hint: our inputs include 0-9
num_epochs = 5    # Number of times we loop through the entire training dataset, can be pretty arbitrary

class NN(nn.Module):
  ############ YOUR CODE STARTS HERE ############
  # 1. Initialize our own NN model
  def __init__(self, input_size, num_classes):
    super(NN, self).__init__()
    self.flatten = nn.Flatten()
    # Use ReLU activation function
    self.relu = nn.ReLU()  
    # Input layer
    self.input_layer = nn.Linear(input_size, 13)  
    # Hidden layers: use at least 1 hidden layer!
    self.hidden1 =  nn.Linear(13, 6) 
    # Output layer
    self.output_layer = nn.Linear(6, 10)
  
  # 2. Define method for forwarding input data
  def forward(self, sample):

    sample = self.flatten(sample)
    out = self.input_layer(sample)  
    out = self.relu(out)  

    out = self.hidden1(out)  
    out = self.relu(out)  #TODO: activation function

    out = self.output_layer(out) #TODO: forward to output layer
    return out
  
nn_model = NN(input_size, num_classes)
print("My NN Model: ", nn_model)

损失函数与优化器 #

loss_function = torch.nn.CrossEntropyLoss()  
optimizer = torch.optim.SGD(nn_model.parameters(), lr=.0001)

训练代码 #

total_steps = len(train_loader)
for epoch in range(num_epochs):
  for i, (images, labels) in enumerate(train_loader):
    # Reshape our images from 2D(28x28) to 1D(784)
    images = images.view(-1, 28*28).to(device)   
    labels = labels.to(device) 

    # Call functions we've previously defined to perform forward pass & calculate loss
    output = nn_model.forward(images) 
    loss = loss_function(output, labels)   

    # Backward pass
    optimizer.zero_grad()  
    loss.backward() 
      # calculates gradients
    optimizer.step() 

    # Print out training process
    if (i+1) % 100 == 0:
      print(f'epoch {epoch+1} / {num_epochs}, step {i+1}/{total_steps}, loss = {loss.item():.4f}')

问题分析与解决建议 #

• 模型容量不足：当前网络隐藏层神经元数量过少（13→6），无法捕捉手写数字的复杂特征。建议扩大网络规模，比如调整为：

  self.input_layer = nn.Linear(input_size, 256)
  self.hidden1 = nn.Linear(256, 128)
  self.output_layer = nn.Linear(128, num_classes)
  

• 学习率设置不合理：SGD优化器的lr=0.0001过小，参数更新幅度不足以让模型向最优解靠近。可以尝试将学习率提升至0.01或0.1，若出现损失震荡再逐步下调；也可以直接改用Adam优化器，其自适应学习率特性更适合新手：

  optimizer = torch.optim.Adam(nn_model.parameters(), lr=0.001)
  

• 训练轮数不足：仅5轮训练不足以让模型充分学习特征，建议将num_epochs调整为15-20轮，观察损失变化趋势。
• 数据未归一化：MNIST图像像素值范围是0-255，未归一化会导致权重更新不稳定。在训练循环中添加数据归一化步骤：

  images = images.view(-1, 28*28).to(device) / 255.0
  

• forward方法调用不规范：PyTorch中标准的模型调用方式是output = nn_model(images)，而非直接调用forward方法，虽然当前功能正常，但建议修正以遵循框架规范。

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