Увеличение потерь обучения после загрузки модели, оптимизатора и планировщика с помощью PyTorch

Question

Я работаю в исследовательском проекте по компьютерному зрению, и у меня есть особая проблема, которая не позволяет мне возобновить обучение после cra sh или прерывания, это мой код для загрузки контрольной точки:

from torch.optim import lr_scheduler
N_EPOCHS = 120
if load_weights:
    optimizer = torch.optim.Adam(model.parameters(),lr=checkpoint['last_lr'], weight_decay=weight_decay)
    optimizer.load_state_dict(checkpoint['optimizer'])
    scheduler = lr_scheduler.StepLR(optimizer, step_size=step_of_scheduler, gamma=0.9, last_epoch=loaded_epochs)
    scheduler.load_state_dict(checkpoint['scheduler'])
else:
    optimizer = torch.optim.Adam(model.parameters(),lr=initial_lr, weight_decay=weight_decay)
    scheduler = lr_scheduler.StepLR(optimizer, step_size=step_of_scheduler, gamma=0.9)

Я даже пытался добавить optimizer.state_dict()['param_groups'][0]['params'] = checkpoint['optimizer']['param_groups'][0]['params'], но результат еще хуже, мой кусок кода для сохранения контрольной точки после проверки:

# Checkpoint
       checkpoint = {'model':model.state_dict(),
                     'epoch':loaded_epochs + epoch + 1,
                     'last_validation_acc': val_acc_db_avg[-1],
                     'hyperparameters': hyperparameters,
                     'last_lr': optimizer.param_groups[0]['lr'],
                     'best_general_val': max(val_acc_db_avg),
                     'last_train_loss': train_loss_db[-1],
                     'optimizer': optimizer.state_dict(),
                     'scheduler': scheduler.state_dict(),
                     'img_resize': 512 if lr_scaled else 256,
                     'best_train_acc': max(train_acc_db),
                     'lr_scaled': lr_scaled
                     }

       # Backup in drive
       torch.save(checkpoint, filename_of_checkpoint)

Я действительно вызвал model.train ( ) до начала тренировки, а также я вызвал scheduler.step (), при этом мое увеличение тренировочной потери составляет от 0,81 до 0,89 или 0,90, если я вручную назначаю параметры, я не знаю, что здесь может быть не так, так как я также правильно загрузите dict состояния модели.

Редактировать: вот как я загружаю dict состояния модели:

if load_weights:
    model.load_state_dict(checkpoint_load['model'])

# Transfering model to GPU if available
model = model.to(DEVICE)