MLflow Trainer Callback

[noise-field]

🚀 Feature request

A callback to log hyperparameters, metrics and cofings/weights to MLFlow, like the existing wandb and Tensorboard callbacks.

Motivation

I use MLFlow as my primary experiment tracking tool. It is convenient to run on a remote server and log the results from any of your training machines, andit also facilitates collaboration.

Trainer is an amazing tool, it makes it very simple to train models, however, the only way to modify the training loop to include custom logging seems to add a callback.

Your contribution

I can contribute a PR.

[julien-c]

Happy to get a PR on this!

[RahulKulhari]

@noise-field is remote tracking with remote server uri with authentication also enabled as part of this feature request ?

[noise-field]

@RahulKulhari well, this is not part of the feature request, but you can certainly do remote tracking with mlflow callback. However, you will need to use environment variables (MLFLOW_TRACKING_URI , MLFLOW_TRACKING_USERNAME, MLFLOW_TRACKING_PASSWORD) in advance to configure your connection to the remote server.