MLflow加载PyTorch Lightning模型时触发ModuleNotFoundError：找不到'src.<mymodelname>'模块

这种问题其实挺常见的——MLflow确实帮你把模型代码打包到了artifacts里，但它不会自动帮你把这些代码路径添加到Python的模块搜索路径中，所以反序列化模型时Python找不到你的src.<mymodelname>模块。下面给你几种解决思路，亲测有效：

核心原因 #

你的PyTorch Lightning模型在序列化（打包）时，pickle会记录模型类的完整模块路径（比如src.mymodelname.MyLightningModel）。当你在外部环境加载模型时，Python需要能从sys.path中找到这个模块的位置，但MLflow默认不会自动把artifacts里的code目录添加到sys.path，所以就触发了ModuleNotFoundError。

解决方法1：手动添加路径（最直接） #

在调用load_model之前，先把artifacts里的code目录加到Python的模块搜索路径里，这样Python就能识别src作为顶层模块了。代码示例如下：

import sys
import os
import mlflow

model_name = "MyModelName"
model_version = "1"
model_uri = f"models:/{model_name}/{model_version}"

# 先把模型artifacts下载到本地临时目录
local_artifacts = mlflow.artifacts.download_artifacts(model_uri=model_uri)
# 找到artifacts里code目录的绝对路径
code_root = os.path.join(local_artifacts, "code")
# 把code目录加到sys.path最前面，确保优先加载这个版本的模型代码（避免和本地其他同名模块冲突）
sys.path.insert(0, code_root)

# 现在再加载模型就没问题了
loaded_model = mlflow.pytorch.load_model(model_uri)

解决方法2：临时添加路径（更安全，避免污染全局环境） #

如果你不想把code目录一直留在sys.path里（怕和其他项目的模块冲突），可以用上下文管理器临时添加路径，加载完模型后自动移除：

import sys
import os
import mlflow
from contextlib import contextmanager

# 定义一个上下文管理器，用于临时添加路径
@contextmanager
def temp_add_sys_path(path):
    sys.path.insert(0, path)
    try:
        yield
    finally:
        # 加载完成后移除路径，恢复环境
        if path in sys.path:
            sys.path.remove(path)

model_name = "MyModelName"
model_version = "1"
model_uri = f"models:/{model_name}/{model_version}"

local_artifacts = mlflow.artifacts.download_artifacts(model_uri=model_uri)
code_root = os.path.join(local_artifacts, "code")

# 临时添加路径，加载模型
with temp_add_sys_path(code_root):
    loaded_model = mlflow.pytorch.load_model(model_uri)

# 离开上下文后，code_root就从sys.path里移除了，不会影响后续代码

关于code_paths的小建议 #

你用code_paths=["src"]把整个src目录打包是合理的，尤其是当模型依赖很多嵌套模块时。不过可以注意两点：

1. 尽量不要打包无关文件（比如测试用例、临时脚本），减小artifacts的体积；
2. 确保模型代码里的import语句和打包后的路径匹配：如果你的模型代码里用的是from src.mymodelname import ...，那就要把code目录加到sys.path；如果用的是相对导入或者直接from mymodelname import ...，那应该把code/src目录加到sys.path。

这样处理后，你加载的模型应该就能正常用于微调或者推理了~

备注：内容来源于stack exchange，提问作者CustardBun