跳到主要内容

Serving 服务编写

目前 OpenBayes 模型部署支持两种部署模式：

  1. 标准的 predictor.py 方式部署
  2. 完全自定义，绕过 OpenBayes 提供的框架的部署

其中自定义的方式是针对需要对部署服务做精细控制的高级用户准备，如果你不确定是不是需要用自定义的方式，那么你就不需要它，建议通过标准方式进行。

标准的 predictor.py 方式​

依赖​

除了业务中用到的库之外，需要额外依赖 openbayes-serving。请尽量保证这个库更新到最新版。

目录结构​

模型部署必须包含两部分内容：

  1. predictor.py 以及其所依赖的文件，用于模型请求的处理
  2. 模型文件

predictor.py 的接口，下文会一一介绍。模型文件的导出在模型导出有详细的介绍。

任何需要在 predictor.py 中引用的文件都需要放同目录或者其子目录下，比如需要一个 classes.json 的文件用于存放分类的信息，则可以以如下方式在 predictor.py 文件中访问：

import json

class Predictor:
    def __init__(self):
        with open('classes.json', 'r') as f:
            values = json.load(f)
        self.values = values

    ...

在 pytorch/image-classifier-resnet50 中包含完整的项目示例。

Predictor​

模板​

Predictor 的模板如下所示：

import openbayes_serving as serv


class Predictor:
    def __init__(self):
        """
        负责加载相应的模型以及对元数据的初始化
        """
        pass

    def predict(self, json):
        """
        在每次请求都会被调用
        接受 HTTP 请求的内容（`json`)
        进行必要的预处理（preprocess）后预测结果，
        最终将结果进行后处理（postprocess）并返回给调用方

        Args:
            json: 请求的数据

        Returns:
            预测结果
        """
        pass

if __name__ == '__main__':  # 如果直接执行了 predictor.py，而不是被其他文件 import
    serv.run(Predictor)  # 开始提供服务

其中 json 参数会依据 HTTP 请求的 Content-Type 头进行解析：

  • 对于 Content-Type: application/json， json 将按照 JSON 格式做解析为字典（Dict）
  • 对于 Content-Type: application/msgpack， 或者其他的 MessagePack 类型的别名，json 会按照与 JSON 格式相同的方式处理将（解析为字典）

样例​

这里提供 pytorch/object-detector 的 predictor.py 文件：

from io import BytesIO

import requests
import torch
from PIL import Image
from torchvision import models
from torchvision import transforms

import openbayes_serving as serv


class Predictor:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"using device: {self.device}")

        model = models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
        model.load_state_dict(torch.load("fasterrcnn_resnet50_fpn_coco-258fb6c6.pth"))
        model = model.to(self.device)
        model.eval()

        self.preprocess = transforms.Compose([transforms.ToTensor()])

        with open("coco_labels.txt") as f:
            self.coco_labels = f.read().splitlines()

        self.model = model

    def predict(self, json):
        # 只使用了 json 参数，其内容为 Dict 格式，直接通过 json[key] 的形式获取内容即可
        threshold = float(json["threshold"])
        image = requests.get(json["url"]).content
        img_pil = Image.open(BytesIO(image))
        img_tensor = self.preprocess(img_pil).to(self.device)
        img_tensor.unsqueeze_(0)

        with torch.no_grad():
            pred = self.model(img_tensor)

        predicted_class = [self.coco_labels[i] for i in pred[0]["labels"].cpu().tolist()]
        predicted_boxes = [
            [(i[0], i[1]), (i[2], i[3])] for i in pred[0]["boxes"].detach().cpu().tolist()
        ]
        predicted_score = pred[0]["scores"].detach().cpu().tolist()
        predicted_t = [predicted_score.index(x) for x in predicted_score if x > threshold]
        if len(predicted_t) == 0:
            return [], []

        predicted_t = predicted_t[-1]
        predicted_boxes = predicted_boxes[: predicted_t + 1]
        predicted_class = predicted_class[: predicted_t + 1]
        return predicted_boxes, predicted_class


if __name__ == '__main__':
    serv.run(Predictor)

Predictor 类的接口​

Predictor 类不需要继承其他的类，但是需要至少提供两个接口：

__init__​

__init__ 的参数个数不定，也没有顺序要求，但是对参数的名字敏感。

每个出现的参数都代表了不同的功能：

  • onnx: 会在 predictor.py 所在的目录搜索 *.onnx 文件加载并作为参数传入

样例：

class PredictorExample1:

    def __init__(self):
        pass

    def predict(self, json):
        return {'result': 'It works!'}
class PredictorExample2:

    def __init__(self, onnx):
        self.onnx = onnx

    def preprocess(self, v):
        ...

    def postprocess(self, v):
        ...

    def predict(self, json):
        onnx = self.onnx
        m = self.preprocess(json['data'])
        result = self.postprocess(onnx.run(None, {'data': m}))
        return result

predict​

predict 的参数个数不定，也没有顺序要求，但是对参数的名字敏感。

每个出现的参数都代表了不同的功能：

  • json: 解析过后的数据，会按照 Content-Type 的值将 JSON 或者 MessagePack 数据解析成 Python 对象。
  • payload: 跟 json 一样，不同的是 json 在遇到无法解析的数据后会返回 400 错误，payload 会直接提供不经解析的数据。
  • data: 不经解析的 POST 数据，永远是 bytes 类型。
  • params: HTTP GET 参数，一个 dict 对象。
  • headers: HTTP 头，一个 dict 对象。
  • request: Flask HTTP Request 对象，请参照 Flask 的文档获取用法。

样例：

class PredictorExample1:

    def __init__(self):
        pass

    def predict(self, json, params):
        return {'message': f'Param foo is {params.get("foo")}'}

predict 函数的返回值​

predict 方法的返回结果可以是以下内容：

  • 一个可以被 JSON 序列化的对象，比如 Python 中的 List, Dict, Tuple
  • 一个 str 类型字符串
  • 一个 bytes 类型对象
  • 一个 flask.Response 类型对象

以下为一些样例：

def predict(self, json):
    return [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
def predict(self):
    return "class 1"
def predict(self):
    # bytes 类型的返回
    array = np.random.randn(3, 3)
    response = pickle.dumps(array)
    return response
def predict(self):
    from flask import Response
    data = b"class 1"
    response = Response(data, mimetype="text/plain")
    return response

完全自定义的方式​

如果部署的模型中存在 start.sh 文件，那么所有的上述的准备工作会全部跳过，直接运行这个文件。 你需要自己完成所有的初始化，以及启动服务的工作。

服务需要监听 80 端口并处理 HTTP 请求。

最后由 Tunghsiao Liu 于 更新