Saving checkpoint, hparams & tfevents after training to separate folder

[dispoth]

Thanks to all the contributors of PyTorch Lightning for a fantastic product!

I want to save a checkpoint, hparams & tfevents after training finishes. I have written this callback:

class AfterTrainCheckpoint(pl.Callback):
    """
    Callback for saving the checkpoint weights, hparams and tf.events after training finishes
    """
    def __init__(self):
        super().__init__()

    def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        print(f"Saving final checkpoint...")
        # As we advance one step at end of training, we use `global_step - 1`
        final_checkpoint_name = f"final_models/final_step_{trainer.global_step - 1}.ckpt"
        final_hparams_name = f"final_models/final_step_{trainer.global_step - 1}.yaml"

        trainer.save_checkpoint(final_checkpoint_name)
        save_hparams_to_yaml(config_yaml=final_hparams_name, hparams=trainer.model.hparams)

1. Is this the best 'Lightning' way to achieve this?
2. How can I save the final events.out.tfevents file to a new directory?
3. Should I be setting save_last=True? seen in ModelCheckpoint in the docs. I am slightly confused about "monitor metrics logged during training/validation steps or end of epochs are not guaranteed to be available at this stage."

[rohitgr7]

hey @dispoth !!

1. I'd say use on_fit_end instead, since the last checkpoint in the model checkpoint is saved in this hook, so it won't guarantee to have that ckpt when your callback calls it.
2. you can copy the log files directly? the are available inside trainer.log_dir.
3. yes, they will be available during both on_train_end and on_fit_end.

[dispoth]

Thanks @rohitgr7 😄

1. Theon_fit_end suggestion looks great! For anyone else interested in the difference between on_fit_end and on_train_end see: https://forums.pytorchlightning.ai/t/are-on-fit-end-and-on-train-end-the-same/412.
2. Copying is fine, I was just hoping for a more 'semantically clear' solution
3. OK perfect