在fastAPI中只加载模型一次

Question

我有一个深度学习模型，我希望在我的API启动时仅加载一次该模型。目前，每次在测试图像的API服务中发出请求时，都会重新加载模型，这需要相当长的时间才能打印输出。 我尝试使用FastAPI的@app.on_event('startup')来加载模型。

main.py
@app.on_event('startup')
def init_data():
    print("init call")
    model = load_model(model_path)
    return model

我希望能将这个模型变量导入到另一个名为classification.py的Python文件中。

#classification.py
pred = model.predict(image)

我不确定这是否是正确的方法。 如有任何帮助，将不胜感激。

谢谢

- NeedToCodeAgain

---

您可以在 main.py 的根级别上加载模型；启动事件的主要用途是为了更轻松地进行测试，并且如果是异步的，则需要在处理任何请求之前完成该函数。在主文件的同步代码中，它应该可以直接使用。另一个选项是在根目录中定义一个字典，然后在启动函数内填充该字典，如参考手册所示：https://fastapi.tiangolo.com/advanced/events/ - MatsLindh

@MatsLindh 我确实尝试了您的第一个建议，即在根级别加载模型，但我无法将在main.py根级别声明的模型变量导入到我的预测文件classification.py中。 - NeedToCodeAgain

Answer 1

所需的只是将其封装在一些函数中，并传递模型即可。

# main.py
import classification
# Other imports...

model = load_model(model_path)
# Model is global so loaded once.


# ...

@app.get('/whatever/page')
def predict(image):
    global model   # We make sure model is not local
    # We call classification's predict function
    result = classification.predict(image, model) 

    #...

#classification.py
# import ...

# The function predict takes the model as a param, and the image to predict.
def predict(image, model):
    pred = model.predict(image)
    # ...
    return pred

你应该更好地输入参数等，但这给出了数据流的想法。另一种方法是将模型加载到 classification.py 的根目录中，这样就不需要传递它。当第一次导入时，它只会被执行一次。

- vinalti

---

我打算遵循FastAPI的建议，在函数参数上使用Python类型提示，例如：def myfunc(txt:str, val:int, arr=List[int])。 - vinalti

嗨，将模型加载到classification.py的根目录后它可以正常工作，谢谢。还有一个问题，是否有办法在我的API服务启动时立即运行scheduler.py文件？现在我需要使用两个终端窗口：1. uvicorn main:app --relaod 和 2. python scheduler.py。我希望我的调度程序文件能够在uvicorn main:app --reload命令执行时立即运行。 - NeedToCodeAgain

你可以创建一个脚本同时启动它们，或者你可以将你的应用程序Docker化。 - vinalti

感谢 @vinalti 的回复，fastapi的start_up解决了我的问题。 - NeedToCodeAgain

@Blackfly 很高兴听到这个消息！你能否在自己的问题答案中描述一下？这对其他人可能会有帮助。 - vinalti