Running test stage after training stage using LightningCLI

[SinaSarparast]

Hello,

I'm using LightningCLI for training a model. I want to run training, validation and testing all in the same experiment to have them all logged in one place, so that they are easily accessible by the logger (mlflow). The validation stage happens after each epoch and the parameters are logged. The testing stage does not take place after training automatically.

I know that I can use fit,validate,test,predict,tune as the documentation says.

$ python trainer.py --help usage: trainer.py [-h] [--config CONFIG] [--print_config [={comments,skip_null}+]] {fit,validate,test,predict,tune} ...

After the training is over the execution of the experiment is terminated and I have manually execute the code and specify the path to the checkpoint. Therefore, a new experiment is created and the test metrics are logged independently as a separate experiment. I want to run the test stage right after training is over with the current model and avoid running a new experiment for testing. I couldn't find a way to do this with the Lightning CLI. How can I run the test stage right after training is over?

Thanks,

[tshu-w]

You can Customize the LightningCLI after_fit, here is an example.
Or implement a callback and add it to LightningCLI by default.