python - mlflow.pytorch.load_model failing due to failure "No module named 'src.<mymodelname>&

admin管理员组

文章数量:1123057

I am working to support some model developers by prototyping some functionality with MLflow model registry.

We successfully register versions of the model fine (it's a Pytorch Lightning model), and I recently also added code_paths to the log_model call so that theoretically all of the original model code ends up registered / in the artifacts (I'm not sure if this is best practice though? But either way the model code is super complex and I'm not sure how I can extract out just the stuff needed by the serializer on the consuming side atm).

        mlflow.pytorch.log_model(
            pytorch_model=cli.model,
            artifact_path="model",
            registered_model_name="MyModelName",
            code_paths=["src"],
        )

In any case, I can see the model registered and the artifacts tree looks like this now when I view them from the UI:

model
|--checkpoints
|--code
    |--src
        |--<mymodelname>
             |-- a bunch of other modules and nested stuff under here
        |--a bunch of other stuff here
|--data
     |--model.pth
     |--pickle_module_info.txt (this is a small file with just the text "mlflow.pytorch.pickle_module" in it)
|--MLmodel  (small file w/ metadata about the model)
|--conda.yaml
|--python_env.yaml
|--requirements.txt

I have a totally separate package outside of the model package and I'm testing consuming the model as an "external" client who wants to use the registered pre-trained model for further fine-tuning or inference. I confirmed that the code is getting retrieved when I do something like this:

model_uri = f"models:/{model_name}/{model_version}"
local_path = mlflow.artifacts.download_artifacts(model_uri)

the retrieved artifacts there include everything I listed above.

However, when I try to load the model:

    loaded_model = mlflow.pytorch.load_model(model_uri)

It fails with the error: ModuleNotFoundError: No module named 'src.<mymodelname>'

Am I going about this totally wrong or is there just something small I'm missing here (like getting the picker to recognize the original model code under code/ somehow since it's all there?)?

Any help would be much appreciated!

本文标签：