Early Stopping Callback restore_best_weights

[PascalIversen]

Hello there!
When I do early stopping, I want to use the best model subsequently.
Reasonable Early Stopping Callbacks therefore should have the restore_best_weights parameter (like in Keras etc.).
Why is this missing from the lightning ES callback? What is the default behavior? If the default is not to restore the weights, how to do it in lightning?
Thanks!

[mpicek]

Hi,

look at ModelCheckpoint callback. After the training, you can use its attribute best_model_path to restore the best model.

Here is a simple code with EarlyStopping and ModelCheckpoint together (training is stopped when val_loss doesn't improve anymore).

# preparing my dataset
dm = MyDataModule()
dm.prepare_data()
dm.setup()

# initialize checkpoints
early_stop = EarlyStopping(monitor="val_loss", patience=10, mode="min")
checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor="val_loss", mode="min")

# initialize the trainer
trainer = pl.Trainer(
    callbacks=[early_stop, checkpoint_callback],   # we use both checkpoints
    max_epochs=100,
    gpus=[0],
    enable_checkpointing=True,
)

# fit the model
trainer.fit(
    model,
    train_dataloaders=dm.train_dataloader(),
    val_dataloaders=dm.val_dataloader(),
)

print(checkpoint_callback.best_model_path)   # prints path to the best model's checkpoint
print(checkpoint_callback.best_model_score) # and prints it score
best_model = MyModel.load_from_checkpoint(checkpoint_callback.best_model_path)

# test only the best model
trainer.test(model=best_model, dataloaders=dm.val_dataloader())