How can I save model weights to mlflow tracking sever using pytorch-lightning?
I would like to save model weights to mlflow tracking using pytorch-lightning.
pytorch-lightning supports logging.
However, it seems that saving model weights as a artifact on mlflow is not supported.

At first, I planed to override ModelCheckpoint class to do it, but I found it is difficult for me because of complex Mixin operations.

Anybody knows simple way to accomplish it?

As @xela said, you can use the experiment object of the mlflow logger to log artifacts.
In case you want to frequently log model weights during training, you could extend ModelCheckpoint:

class MLFlowModelCheckpoint(ModelCheckpoint):
    def __init__(self, mlflow_logger, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mlflow_logger = mlflow_logger

    @rank_zero_only
    def on_validation_end(self, trainer, pl_module):
        super().on_validation_end(trainer, pl_module)
        run_id = self.mlflow_logger.run_id
        self.mlflow_logger.experiment.log_artifact(run_id, self.best_model_path)

And then use in your training code
mlflow_logger = MLFlowLogger()
checkpoint_callback = MLFlowModelCheckpoint(mlflow_logger)
trainer = pl.Trainer(checkpoint_callback=checkpoint_callback, logger=mlflow_logger)

