MLFlowLogger implementation: checkpoints downloaded during mlflow.pytorch.load_model.

[croth1-liveeo]

TL;DR: It appears to me that lightnings MLFlow logger implementation's choice to put checkpoints into mlflow's model directory might lead to excessive downloads during model loading.

Details

It is my understanding that checkpoints are generally much larger than the models themselves, probably due to extra information such as optimizer state, etc.

Lightning's MLflow logger implementation

Furthermore, it seems that the default behavior of lightnings mlflowlogger is to upload checkpoints at the end of the run according to the ModelCheckpoints config. The checkpoint path seems hard-coded to the checkpoints subdirectory of the MLflow model.

pytorch-lightning/src/lightning/pytorch/loggers/mlflow.py

Line 356 in 8646515

artifact_path = f"model/checkpoints/{Path(p).stem}"

MLflows model loading behavior

the MLflow model loading code seems to download the full models directory when loading from a remote target.

To me this seems to mean that whenever one attempts to load such a model with mlflow.pytorch.load_model(f"runs:/{run_id}/model") logged with lightnings mlflow-logger implementation, one would always also download all the checkpoints. However it is my suspicion that the weights needed during unpickling are in data/model.pth.

Moving the checkpoints out of the model folder?

Ultimately my question is: is there a reason for putting the checkpoints as a sub-directory into model directory, or would it be better to put them under a top-level checkpoints directory which to my understanding should avoid downloading all the checkpoints during model load?

Am I missing something here?

[croth1-liveeo]

Seems to have been us using the model directory for logging the model explicitly, clashing with lightnings choice of using model/checkpoints as their preferred checkpoint location.

Bit of a gotcha I was not foreseeing without looking deep into the code. Is there a reason for choosing model/checkpoints instead of letting the user set the artifact location?