How to recover Best Epoch from Checkpoint/Earlystop hook?

[thiagoribeirodamotta]

Hello there,

I've recently created a script to try out different hyperperameters when training. This script will change some core aspects of the input config file programatically and then proceed to perform a trainning run with the new modifications.

So far everything is working as expected, but in order to better log the difference between runs, I decided to use MLFlow.

Since MLFlow is pretty new to OpenMMLab, there isn't a tutorial on how to use it yet, but according to the documentation, this is what I did:

for model_name in model_names:
   for model_arch in model_arches:

        vis_backend = MLflowVisBackend(save_dir='./mlruns/', 
                                                       tracking_uri="http://127.0.0.1:5000",
                                                       exp_name="BaseConfigFactory",
                                                       run_name=model_name + "_" + model_arch + "_" + finetune_name + "_" + time.strftime("%Y_%m%d_%H%M"),
                                                       tags = {'User': 'my_user'})
        vis_backend._init_env()                
        vis_backend.cfg = updated_cfg
        
        # ... fill logged_params as a dict
        
        print("Logging parameters and emptying cuda cache...")
        [vis_backend._mlflow.log_param(k, v) for k, v in logged_params.items()]
        torch.cuda.empty_cache()

        runner, model = train(updated_cfg, amp=True, work_dir=updated_cfg.work_dir)
        val_metrics = runner.val()
        test_metrics = runner.test()

        val_metrics = {k.replace('/', '/val/'): v for k, v in val_metrics.items()}
        test_metrics = {k.replace('/', '/test/'): v for k, v in test_metrics.items()}
        
        vis_backend.add_scalars(val_metrics)
        vis_backend.add_scalars(test_metrics)
        vis_backend.close()
        print("Cleaning model and runner...")
        del model
        del runner
        del vis_backend
        print("Emptying cache cuda...")
        torch.cuda.empty_cache()

With this, I can successfully log most of the parameters I want. But now I'd like to log the epoch which resulted in the best model saved by the checkpoint/earlystop hook.

How can I recover the best epoch from those hooks?