machine-learning - 从检查点恢复时的训练损失爆炸式增长

Question

我正在尝试在我的算法中实现一个函数，该函数允许我从检查点恢复训练。问题是，当我恢复训练时，我的损失会爆炸许多数量级，从 0.001 到 1000 的数量级。我怀疑问题可能是恢复训练时，没有正确设置学习率。

这是我的训练功能：

def train_gray(epoch, data_loader, device, model, criterion, optimizer, i, path):
    train_loss = 0.0
    for data in data_loader:
        img, _ = data
        img = img.to(device)
        stand_dev = 0.0392

        noisy_img = add_noise(img, stand_dev, device)

        output = model(noisy_img, stand_dev)
        output = output[:,0:1,:,:]  
        loss = criterion(output, img)

        optimizer.zero_grad()

        loss.backward()
        optimizer.step()
        train_loss += loss.item()*img.size(0)

    train_loss = train_loss/len(data_loader)

    print('Epoch: {} Complete \tTraining Loss: {:.6f}'.format(
        epoch,
        train_loss
        ))
    return train_loss

这是我的主要函数，它初始化我的变量，加载一个检查点，调用我的训练函数，并在一个训练周期后保存一个检查点：

def main():
    now = datetime.now()
    current_time = now.strftime("%H_%M_%S")
    path = "/home/bledc/my_remote_folder/denoiser/models/{}_sigma_10_session2".format(current_time)
    os.mkdir(path)

    width = 256
    # height = 256
    num_epochs = 25
    batch_size = 4
    learning_rate = 0.0001

    data_loader = load_dataset(batch_size, width)
    
    model = UNetWithResnet50Encoder().to(device)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=learning_rate, weight_decay=1e-5)

    ############################################################################################
    # UNCOMMENT CODE BELOW TO RESUME TRAINING FROM A MODEL
    model_path = "/home/bledc/my_remote_folder/denoiser/models/resnet_sigma_10/model_epoch_10.pt"
    save_point = torch.load(model_path)
    model.load_state_dict(save_point['model_state_dict'])
    optimizer.load_state_dict(save_point['optimizer_state_dict'])
    epoch = save_point['epoch']
    train_loss = save_point['train_loss']
    model.train()
    ############################################################################################

    for i in range(epoch, num_epochs+1):
        train_loss = train_gray(i, data_loader, device, model, criterion, optimizer, i, path)
        checkpoint(i, train_loss, model, optimizer, path)

    print("end")

最后，这是我保存检查点的功能：

def checkpoint(epoch, train_loss, model, optimizer, path):
    torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'train_loss': train_loss
            }, path+"/model_epoch_{}.pt".format(epoch))
    print("Epoch saved")

如果我的问题是我没有保存我的学习率，我该怎么做？

任何帮助将不胜感激，克莱门特

更新：我相当确定问题出在我的预训练模型上。我在每个时期都保存优化器，但优化器只保存可训练层的信息。我希望尽快解决这个问题，并在我确定谁来保存和加载整个模型时发布更彻底的答案。

标签： machine-learningdeep-learningpytorchtorchvision