自定义数据集训练简单Pytorch模型首个epoch后停滞问题求助

在自行创建的程序化Dataset上训练简单PyTorch模型时，训练过程在首个epoch结束后完全停滞。数据集已完成归一化、打乱操作，单batch数据输入未训练模型可正常得到输出，仅训练环节存在异常。

相关代码 #

批次生成代码 #

import math

def makebatches(x_train,y_train,batchsize):
  numBatches = math.ceil(len(x_train)/batchsize)
  batches = []
  for i in range(numBatches):
    batches.append((x_train[i*batchsize:(i+1)*batchsize],y_train[i*batchsize:(i+1)*batchsize]))
  return batches

训练代码 #

BATCH_SIZE = 32
batches = makebatches(images,labels,BATCH_SIZE)
learning_rate = 0.01
num_epochs = 10

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyModel()
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer =  torch.optim.Adam(model.parameters(), lr = 0.01)  

for epoch in range(num_epochs):
    train_running_loss = 0.0
    train_acc = 0.0

    model = model.train()

    ## training step
    for i, (X, Y) in enumerate(batches):
        
        X = X.to(device)
        Y = Y.to(device)

        ## forward + backprop + loss
        logits = model(X)
        loss = criterion(logits, Y)
        optimizer.zero_grad()
        loss.backward()

        ## update model params
        optimizer.step()

        train_running_loss += loss.detach().item()
        train_acc += get_accuracy(logits, Y, BATCH_SIZE)
        #print(i)
    
    model.eval()
    print('Epoch: %d | Loss: %.4f | Train Accuracy: %.2f' \
          %(epoch, train_running_loss / i, train_acc/i)) 

训练输出结果 #

Epoch: 0 | Loss: 0.8300 | Train Accuracy: 53.26
Epoch: 1 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 2 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 3 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 4 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 5 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 6 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 7 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 8 | Loss: 0.8445 | Train Accuracy: 52.34
Epoch: 9 | Loss: 0.8445 | Train Accuracy: 52.34

1. 固定批次无随机打乱：批次在训练循环外一次性生成，所有epoch都使用完全相同顺序的批次，模型快速过拟合到固定数据顺序，陷入局部最优后不再更新。
2. 学习率设置过高：Adam优化器的常规适用学习率为1e-3~1e-4，代码中设置的0.01超出合理区间，过大的学习率会导致参数更新幅度过大，直接跳过最优解，收敛到次优平坦区域后停滞。
3. 指标计算存在偏差：计算平均损失和准确率时使用i作为除数，enumerate索引从0开始，遍历完所有批次后i等于总批次数量减1，计算出的平均值存在误差，会干扰对训练状态的判断。

• 每个 epoch 重新生成打乱后的批次：取消训练循环外的批次生成逻辑，每个epoch开始前先将特征和标签同步打乱，再切分批次，保证每轮训练的数据顺序不同。
• 调低优化器学习率：将Adam的学习率调整为1e-3，即optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)，匹配该优化器的常规适用区间。
• 修正平均指标计算逻辑：提前获取总批次数量num_batches = len(batches)，输出时用总批次数量作为除数计算平均损失和准确率，避免索引偏差导致的计算错误。
• 推荐替换自定义批次逻辑：直接使用PyTorch官方的DataLoader实现批次生成和打乱，官方实现经过优化，稳定性更高，示例代码如下：

from torch.utils.data import TensorDataset, DataLoader

# 封装数据集
train_dataset = TensorDataset(images, labels)
# 生成迭代器，自动完成打乱、切分批次操作
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

训练时直接遍历train_loader即可，不需要自行实现makebatches函数。

内容的提问来源于stack exchange，提问作者José Roberto Canuto Vasconcelo