Ошибка mlflow.pytorch.load_model из-за сбоя «Нет модуля с именем 'src.'» с unpickler.load ⇐ Python

[Anonymous]

Я работаю над поддержкой некоторых разработчиков моделей, создавая прототипы некоторых функций с помощью реестра моделей MLflow.
Мы успешно регистрируем версии модели (это модель Pytorch Lightning), и недавно я также добавил code_paths в вызов log_model, чтобы теоретически весь исходный код модели был зарегистрирован/в артефактах (хотя я не уверен, что это лучшая практика? Но в любом случае код модели очень сложен, и я не конечно, как я может извлечь только то, что необходимо сериализатору на стороне потребления).

Код: Выделить всё

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

В любом случае я вижу зарегистрированную модель и дерево артефактов, когда я просматриваю их из пользовательского интерфейса:

Код: Выделить всё

model
|--checkpoints
|--code
|--src
|--
|-- 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

У меня есть совершенно отдельный пакет помимо пакета модели, и я тестирую использование модели в качестве «внешнего» клиента, который хочет использовать зарегистрированную предварительно обученную модель для дальнейшей доработки. настройка или вывод. Я подтвердил, что код извлекается, когда я делаю что-то вроде этого:

Код: Выделить всё

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

Извлеченные артефакты включают в себя все, что я перечислил выше.
Однако, когда я пытаюсь загрузить модель:

Код: Выделить всё

    loaded_model = mlflow.pytorch.load_model(model_uri)

Это завершается с ошибкой: ModuleNotFoundError: Нет модуля с именем 'src.'
Я собираюсь сделать это полностью неправильно, или мне здесь не хватает чего-то маленького (например, заставить средство выбора каким-то образом распознать исходный код модели под кодом, поскольку все это есть?)?
Любая помощь будет очень полезна оценил!

Подробнее здесь: https://stackoverflow.com/questions/793 ... c-mymodeln