使用预训练模型的 FastAPI 预测

使用 FastAPI 在生产环境中部署 AI

Matt Eckerle

Software and Data Engineering Leader

环境设置

所需库：

FastAPI：用于构建 Python API 的框架
uvicorn：用于运行 Python Web 应用的高速 ASGI 服务器
joblib：用于加载模型

from fastapi import FastAPI
import uvicorn
import joblib
# Create the FastAPI app instance
app = FastAPI()

加载预训练的企鹅分类器

基于 Palmer Penguins 数据集训练
根据 4 个特征预测企鹅物种：喙长、喙深、鳍长、体重
输出：Adelie、Chinstrap 或 Gentoo

import joblib

# Load the pre-trained model
model = joblib.load('penguin_classifier.pkl')
# Check data type of model to verify model loading
print(type(model))

<class 'sklearn.pipeline.Pipeline'>

¹ https://huggingface.co/SIH/penguin-classifier-sklearn

Uvicorn

ASGI（异步服务器网关接口）服务器
由 Python 构建并服务于 Python

uvicorn main:app \
        --host 0.0.0.0 \
        --port 8080

import uvicorn
uvicorn.run(app, 
            host="0.0.0.0", 
            port=8080)

Uvicorn 标志

创建预测端点

# FastAPI prediction endpoint
@app.post("/predict")
def predict(culmen_length_mm, culmen_depth_mm, 
            flipper_length_mm, body_mass_g):

    features = [[culmen_length_mm, culmen_depth_mm,
                 flipper_length_mm, body_mass_g]]

    prediction = model.predict(features)[0]
    return {"predicted_species": prediction}

运行应用程序

if __name__ == "__main__":
    uvicorn.run(
      app, 
      host="0.0.0.0", 
      port=8080)

将全部代码保存为 Python 文件：your_api_script.py

$ python3 your_api_script.py

Uvicorn 启动日志

测试 API

curl \

  -X POST "http://localhost:8080/predict" \

  -H "Content-Type: application/json" \

  -d '{"culmen_length_mm": 39.1,
       "culmen_depth_mm": 18.7,
       "flipper_length_mm": 181,
       "body_mass_g": 3750}'

{
    "prediction": "Adelie",
    "confidence": 0.87
}

Passons à la pratique !

使用 FastAPI 在生产环境中部署 AI