Блог

Я пытаюсь реализовать функцию в своем алгоритме, которая позволяет мне возобновить обучение с контрольной точки. Проблема в том, что когда я возобновляю обучение, мои потери увеличиваются на много порядков, от порядка 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")
  

Если моя проблема в том, что я не сохраняю скорость обучения, как мне это сделать?

Любая помощь будет с благодарностью, Клемент

Обновление: я совершенно уверен, что проблема заключается в моей предварительно обученной модели. Я сохраняю оптимизатор каждую эпоху, но оптимизатор хранит информацию только для обучаемых слоев. Я надеюсь решить эту проблему в ближайшее время и опубликовать более подробный ответ, когда выясню, кого сохранять и загружать всю модель.