PyTest运行加载FastAI模型时出现AttributeError：找不到my_func

问题原因 #

直接用python3运行测试脚本时，脚本以__main__模块身份执行，你导入的my_func存在于这个__main__命名空间中，因此FastAI加载模型时（底层用torch pickle反序列化）能找到该函数。

但用pytest运行时，测试脚本会被作为pytest.__main__模块执行，你导入的my_func实际存在于tests.test_predict模块命名空间，而非pytest.__main__，导致pickle反序列化时找不到函数，抛出AttributeError。

解决方案 #

方案1：将自定义函数注入pytest的__main__命名空间 #

在测试脚本开头，导入my_func后，显式把它添加到pytest.__main__的模块命名空间中：

import pytest
import sys
from fastai.vision.all import *
from project.predict import my_func, prediction

# 把my_func注入到pytest.__main__命名空间
sys.modules['pytest.__main__'].my_func = my_func

model = load_learner('model_path/', cpu=True)
model.load(fast_ai_params.weights)

def my_test():
    pred = prediction(tile_params, model, full_tile_path)
    assert pred.shape

方案2：用Pytest Fixture延迟模型加载（推荐） #

将模型加载逻辑放到Pytest Fixture中，避免在模块级别执行，同时确保函数所在模块的命名空间被正确识别：

import pytest
import sys
from fastai.vision.all import *
from project.predict import my_func, prediction

# 确保my_func所在的模块被pickle识别
import project.predict
sys.modules['project.predict'] = project.predict

@pytest.fixture(scope="module")
def trained_model():
    learner = load_learner('model_path/', cpu=True)
    learner.load(fast_ai_params.weights)
    return learner

def my_test(trained_model):
    pred = prediction(tile_params, trained_model, full_tile_path)
    assert pred.shape

这种方式不仅解决了命名空间问题，还能利用Fixture的作用域复用模型实例，提升测试效率。

方案3：修改模型保存时的函数引用（根源解决） #

如果有权限重新训练保存模型，可以在训练时确保my_func被绑定到固定模块（而非__main__）。比如在训练脚本中，不要在__main__里定义my_func，而是把它放到project.predict模块中，然后从该模块导入使用，再保存模型。这样后续加载时，pickle会直接去project.predict模块找函数，不会受运行环境的命名空间影响。

内容的提问来源于stack exchange，提问作者DanielBell99