发布于 由admin

Python模型部署：利用Flask、FastAPI和ONNX Runtime将机器学习模型部署为API服务。

Python模型部署：Flask、FastAPI与ONNX Runtime实战

大家好！今天我们来聊聊如何将我们辛辛苦苦训练好的机器学习模型部署为API服务，让其他人也能方便地使用它。我们将重点介绍三种主流的技术栈：Flask、FastAPI和ONNX Runtime，并结合实际代码进行讲解。

一、模型部署的必要性与流程

在机器学习项目的生命周期中，模型训练仅仅是其中一步。更重要的是如何将训练好的模型应用到实际场景中，为用户提供服务。这就是模型部署的意义所在。

模型部署的主要流程通常包括以下几个步骤：

  1. 模型训练与评估： 这是基础，我们需要得到一个性能良好的模型。
  2. 模型序列化： 将训练好的模型保存到磁盘，方便后续加载。
  3. API服务构建： 使用Web框架（如Flask或FastAPI）搭建API接口，接收用户请求并返回预测结果。
  4. 模型加载与推理： 在API服务中加载模型，对接收到的数据进行预处理，然后进行推理，得到预测结果。
  5. 部署与监控： 将API服务部署到服务器，并进行监控，确保服务稳定运行。

二、Flask：轻量级Web框架入门

Flask是一个轻量级的Python Web框架，简单易用，非常适合快速构建API服务。

1. 安装Flask：

pip install Flask

2. 一个简单的Flask应用：

from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json(force=True)  # 获取POST请求中的JSON数据
    # 假设这里有一个predict函数，用于进行预测
    prediction = predict_function(data)
    return jsonify(prediction)  # 返回JSON格式的预测结果

def predict_function(data):
    # 这里模拟一个简单的预测函数
    input_value = data['input']
    result = input_value * 2  # 简单的乘以2
    return {'result': result}

if __name__ == '__main__':
    app.run(debug=True)

代码解释：

  • from flask import Flask, request, jsonify: 导入Flask框架以及request和jsonify模块，分别用于处理请求和返回JSON数据。
  • app = Flask(__name__): 创建一个Flask应用实例。
  • @app.route('/predict', methods=['POST']): 定义一个路由，当接收到/predict路径的POST请求时，执行predict函数。
  • request.get_json(force=True): 获取POST请求中的JSON数据。force=True表示强制将请求体解析为JSON格式。
  • jsonify(prediction): 将预测结果转换为JSON格式，并返回给客户端。
  • predict_function(data): 这是一个占位符，实际应用中需要替换成你的机器学习模型的预测函数。
  • app.run(debug=True): 启动Flask应用。debug=True表示开启调试模式，方便开发调试。

3. 模型加载与推理（以sklearn模型为例）：

假设我们有一个使用sklearn训练好的线性回归模型linear_model.pkl。

import pickle
from flask import Flask, request, jsonify

app = Flask(__name__)

# 加载模型
with open('linear_model.pkl', 'rb') as f:
    model = pickle.load(f)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json(force=True)
    try:
        # 从JSON数据中提取特征
        features = [data['feature1'], data['feature2']] # 假设模型需要两个特征
        prediction = model.predict([features])[0] # 注意这里要传入一个二维数组，并取第一个元素
        return jsonify({'prediction': prediction})
    except Exception as e:
        return jsonify({'error': str(e)}), 400 # 返回错误信息和HTTP状态码400

if __name__ == '__main__':
    app.run(debug=True)

代码解释：

  • import pickle: 导入pickle模块，用于加载模型。
  • with open('linear_model.pkl', 'rb') as f: model = pickle.load(f): 使用pickle加载模型。
  • features = [data['feature1'], data['feature2']]: 从JSON数据中提取特征。请根据你的模型实际需要的特征进行修改。
  • prediction = model.predict([features])[0]: 使用模型进行预测。注意sklearn的predict方法需要传入一个二维数组，即使只有一个样本，也要将其转换为二维数组。[0]是为了提取预测结果的第一个元素，因为predict方法返回的是一个数组。
  • return jsonify({'error': str(e)}), 400: 如果发生异常，返回错误信息和HTTP状态码400，表示客户端请求错误。

4. 运行Flask应用：

保存以上代码为app.py，然后在命令行中运行：

python app.py

5. 测试API：

可以使用curl或者Postman等工具发送POST请求到http://127.0.0.1:5000/predict，请求体为JSON格式的数据：

{
  "feature1": 1.0,
  "feature2": 2.0
}

如果一切正常，你将会收到包含预测结果的JSON响应。

三、FastAPI：高性能异步Web框架

FastAPI是一个现代的、高性能的Python Web框架，主要用于构建API。它具有以下优点：

  • 速度快： 基于Starlette和Pydantic，性能非常出色。
  • 易于使用： 提供了自动化的数据验证、序列化和文档生成。
  • 类型提示： 强制使用类型提示，提高代码可读性和可维护性。

1. 安装FastAPI：

pip install fastapi uvicorn

2. 一个简单的FastAPI应用：

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI()

class InputData(BaseModel):
    input: float

class OutputData(BaseModel):
    result: float

@app.post("/predict", response_model=OutputData)
async def predict(data: InputData):
    # 假设这里有一个predict函数，用于进行预测
    prediction = predict_function(data.input)
    return {"result": prediction}

def predict_function(input_value: float) -> float:
    # 这里模拟一个简单的预测函数
    result = input_value * 2  # 简单的乘以2
    return result

代码解释：

  • from fastapi import FastAPI, HTTPException: 导入FastAPI框架以及HTTPException模块，用于处理异常。
  • from pydantic import BaseModel: 导入Pydantic的BaseModel，用于数据验证和序列化。
  • app = FastAPI(): 创建一个FastAPI应用实例。
  • class InputData(BaseModel): input: float: 定义输入数据模型，使用Pydantic进行数据验证。input: float表示输入数据必须是一个浮点数。
  • class OutputData(BaseModel): result: float: 定义输出数据模型，使用Pydantic进行数据验证。result: float表示输出数据必须是一个浮点数。
  • @app.post("/predict", response_model=OutputData): 定义一个POST路由，当接收到/predict路径的POST请求时，执行predict函数。response_model=OutputData表示返回的数据必须符合OutputData模型的定义。
  • async def predict(data: InputData): 定义一个异步函数，用于处理预测请求。data: InputData表示输入数据必须符合InputData模型的定义。
  • return {"result": prediction}: 返回预测结果，FastAPI会自动将其序列化为JSON格式。

3. 模型加载与推理（以sklearn模型为例）：

import pickle
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI()

# 加载模型
with open('linear_model.pkl', 'rb') as f:
    model = pickle.load(f)

class InputData(BaseModel):
    feature1: float
    feature2: float

class OutputData(BaseModel):
    prediction: float

@app.post("/predict", response_model=OutputData)
async def predict(data: InputData):
    try:
        features = [data.feature1, data.feature2]
        prediction = model.predict([features])[0]
        return {"prediction": prediction}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

代码解释：

  • class InputData(BaseModel): feature1: float; feature2: float: 定义输入数据模型，包含两个浮点数类型的特征。
  • raise HTTPException(status_code=400, detail=str(e)): 如果发生异常，抛出一个HTTPException异常，FastAPI会自动将其转换为HTTP响应，并返回错误信息和状态码。

4. 运行FastAPI应用：

保存以上代码为main.py，然后在命令行中运行：

uvicorn main:app --reload

--reload参数表示开启自动重载，当代码发生修改时，服务会自动重启。

5. 测试API：

可以使用curl或者Postman等工具发送POST请求到http://127.0.0.1:8000/predict，请求体为JSON格式的数据：

{
  "feature1": 1.0,
  "feature2": 2.0
}

FastAPI还提供了自动化的API文档，可以通过访问http://127.0.0.1:8000/docs查看。

四、ONNX Runtime：跨平台高性能推理引擎

ONNX (Open Neural Network Exchange) 是一种开放的模型表示格式，允许在不同的深度学习框架之间迁移模型。ONNX Runtime是一个跨平台的高性能推理引擎，可以加速ONNX模型的推理速度。

1. ONNX的优势：

  • 跨平台： 可以在不同的操作系统和硬件平台上运行。
  • 高性能： 对主流硬件平台进行了优化，可以加速模型推理。
  • 框架无关性： 可以在不同的深度学习框架之间迁移模型。

2. 将模型转换为ONNX格式（以PyTorch为例）：

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(2, 1)  # 输入维度为2，输出维度为1

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = SimpleModel()

# 创建一些随机输入数据
dummy_input = torch.randn(1, 2) # 批次大小为1，输入维度为2

# 导出模型为ONNX格式
torch.onnx.export(model,
                  dummy_input,
                  "simple_model.onnx",
                  verbose=True,
                  input_names = ['input'],   # 输入节点的名称
                  output_names = ['output'], # 输出节点的名称
                  dynamic_axes={'input' : {0 : 'batch_size'},    # 动态批次大小
                                'output' : {0 : 'batch_size'}})

代码解释：

  • import torch: 导入PyTorch库。
  • class SimpleModel(nn.Module): 定义一个简单的PyTorch模型。
  • dummy_input = torch.randn(1, 2): 创建一个随机输入数据，用于导出ONNX模型。
  • torch.onnx.export(...): 使用torch.onnx.export函数将PyTorch模型导出为ONNX格式。
    • model: 要导出的模型实例。
    • dummy_input: 一个示例输入数据，用于确定模型的输入和输出形状。
    • "simple_model.onnx": 导出的ONNX模型的文件名。
    • verbose=True: 是否打印导出过程的详细信息。
    • input_names: 输入节点的名称，方便后续使用ONNX Runtime进行推理。
    • output_names: 输出节点的名称，方便后续使用ONNX Runtime进行推理。
    • dynamic_axes: 定义动态轴，允许模型处理不同大小的输入。这里将批次大小设置为动态的。

3. 安装ONNX Runtime：

pip install onnxruntime

4. 使用ONNX Runtime进行推理：

import onnxruntime
import numpy as np

# 加载ONNX模型
sess = onnxruntime.InferenceSession("simple_model.onnx")

# 获取输入和输出节点的名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 创建输入数据
input_data = np.array([[1.0, 2.0]], dtype=np.float32) # 注意数据类型需要匹配模型的要求

# 进行推理
output = sess.run([output_name], {input_name: input_data})[0]

# 打印结果
print(output)

代码解释：

  • import onnxruntime: 导入ONNX Runtime库。
  • sess = onnxruntime.InferenceSession("simple_model.onnx"): 加载ONNX模型。
  • input_name = sess.get_inputs()[0].name: 获取输入节点的名称。
  • output_name = sess.get_outputs()[0].name: 获取输出节点的名称。
  • input_data = np.array([[1.0, 2.0]], dtype=np.float32): 创建输入数据，注意数据类型需要匹配模型的要求。
  • output = sess.run([output_name], {input_name: input_data})[0]: 使用ONNX Runtime进行推理。
    • [output_name]: 指定要获取的输出节点的名称。
    • {input_name: input_data}: 传入输入数据。
    • [0]: 提取预测结果。

5. 将ONNX Runtime集成到Flask或FastAPI中：

这里以FastAPI为例：

import onnxruntime
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI()

# 加载ONNX模型
sess = onnxruntime.InferenceSession("simple_model.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

class InputData(BaseModel):
    feature1: float
    feature2: float

class OutputData(BaseModel):
    prediction: float

@app.post("/predict", response_model=OutputData)
async def predict(data: InputData):
    try:
        input_data = np.array([[data.feature1, data.feature2]], dtype=np.float32)
        output = sess.run([output_name], {input_name: input_data})[0]
        prediction = output[0][0] # 提取预测结果
        return {"prediction": prediction}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

代码解释：

  • 在应用启动时加载ONNX模型，避免每次请求都加载模型，提高性能。
  • 将输入数据转换为NumPy数组，并指定数据类型为np.float32，确保数据类型与模型要求一致。
  • 提取预测结果，注意output是一个二维数组，需要提取第一个元素的第一个元素。

五、模型部署的注意事项

  • 模型版本控制： 使用Git等工具进行模型版本控制，方便回滚和管理。
  • 性能监控： 监控API服务的性能指标，如请求延迟、吞吐量等，及时发现和解决问题。
  • 安全性： 对API服务进行安全加固，防止恶意攻击。
  • 日志记录： 记录API服务的日志，方便排查问题。
  • 错误处理： 完善的错误处理机制，能够返回清晰的错误信息，方便客户端调试。
  • 数据验证： 对输入数据进行验证，防止恶意输入。
  • 资源管理： 合理分配服务器资源，避免资源耗尽。
  • 模型更新： 提供模型更新机制，方便更新模型。

六、总结：部署模型，服务大众

我们学习了如何使用Flask、FastAPI和ONNX Runtime将机器学习模型部署为API服务。根据项目的实际情况选择合适的技术栈，并注意模型部署的各项注意事项，才能构建稳定、高效的API服务。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注