FastAPI将机器学习模型转换为Python API示例

学习如何使用 FastAPI 库将机器学习模型转换为 Python 中的 API。FastAPI 是一个现代、快速（高性能）的 Web 框架，用于构建 API，它基于标准 Python 类型提示。我们将使用Titanic 数据集和逻辑回归模型（示例）。

为什么选择 FastAPI？

FastAPI 具有以下优点：

• 基于 Python 类型提示，自动生成文档。
• 性能优异，与 Flask 相比，FastAPI 更快。
• 支持异步请求处理。

环境准备

首先，确保您的环境中安装了以下库：

• fastapi
• uvicorn
• pandas
• scikit-learn
• numpy

您可以使用 pip 安装这些库：

pip install fastapi uvicorn pandas scikit-learn numpy

步骤 1: 加载和预处理数据

我们将使用 Titanic 数据集，并进行必要的预处理。

import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.impute import SimpleImputer

# 加载数据
df = pd.read_csv("train.csv")
include = ['Age', 'Sex', 'Embarked', 'Survived']
df = df[include]

# 处理缺失值
imputer = SimpleImputer(strategy='median')
df['Age'] = imputer.fit_transform(df[['Age']])

# 编码分类特征
label_encoder = LabelEncoder()
df['Sex'] = label_encoder.fit_transform(df['Sex'])

onehot_encoder = OneHotEncoder(sparse=False, drop='first', handle_unknown='ignore')
df = pd.get_dummies(df, columns=['Embarked'], dummy_na=True)

步骤 2: 训练机器学习模型

我们将使用逻辑回归模型进行训练。

from sklearn.linear_model import LogisticRegression

# 定义特征和标签
X = df.drop('Survived', axis=1)
y = df['Survived']

# 训练模型
model = LogisticRegression()
model.fit(X, y)

步骤 3: 创建 FastAPI 应用

我们将创建一个 FastAPI 应用，用于加载模型并提供预测。

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sklearn.externals import joblib

app = FastAPI()

# 加载模型
model = joblib.load('model.pkl')  # 需要将前面的模型保存，命名model.pkl

class PredictionRequest(BaseModel):
    Age: float
    Sex: int
    Embarked_C: int = 0
    Embarked_Q: int = 0
    Embarked_S: int = 0
    Embarked_nan: int = 0

@app.post("/predict/")
async def predict(request: PredictionRequest):
    """
    接收乘客信息并返回生存预测。
    """
    try:
        # 创建 DataFrame 并进行预测
        input_df = pd.DataFrame([request.dict()])
        prediction = model.predict(input_df)
        return {"prediction": int(prediction[0])}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

步骤 4: 运行 FastAPI 应用

使用 Uvicorn 作为 ASGI 服务器运行 FastAPI 应用。

uvicorn your_script_name:app --reload

将 your_script_name 替换为您的 Python 脚本名称。

测试 API

您可以使用 Postman 或任何其他 API 客户端测试您的 API。发送 POST 请求到 http://127.0.0.1:8000/predict/ 并包含 JSON 负载，例如：

{
  "Age": 25.0,
  "Sex": 1,
  "Embarked_C": 0,
  "Embarked_Q": 0,
  "Embarked_S": 1,
  "Embarked_nan": 0
}

总结

在本教程中，我们学习了如何使用 FastAPI 将机器学习模型转换为 Python 中的 API。我们预处理了数据，训练了一个逻辑回归模型，并创建了一个简单的 FastAPI 应用来提供预测。FastAPI 提供了一种快速且高效的方式来构建和部署机器学习模型的 API。

来源：python收藏家