Python调用TensorFlow Serving部署的模型提供接口

安装Docker

参考：Ubuntu 16.04安装docker详细步骤

Pull TensorFlow Serving

docker pull tensorflow/serving

PyTorch或TensorFlow模型生成模型文件

模型生成.onnx文件

Pytorch模型生成.onnx文件

import torch

def torch2onnx(device: str, onnx_path: str):
    ...
    model.eval()
    input_names = ['image']
    for k, v in model.named_parameters():
        input_names.append(k)
        v.requires_grad = False

    dummy_input = torch.rand(1, 3, 640, 640).to(device)

    output_names = ['output']
    torch.onnx.export(model, dummy_input, onnx_path,
                      verbose=True,
                      input_names=input_names,
                      output_names=output_names,
                      opset_version=11)


if __name__ == '__main__':
    torch2onnx(device='cpu', onnx_path='./model.onnx')

.h5文件转.onnx文件

import keras2onnx
import onnx
from keras.models import load_model
model = load_model('./model.h5')
onnx_model = keras2onnx.convert_keras(model, model.name)
temp_model_file = './model.onnx'
onnx.save_model(onnx_model, temp_model_file)

.onnx文件转.pb文件

import onnx
from onnx_tf import backend


def onnx2tf(onnx_path: str, tf_dir: str):
    onnx_model = onnx.load(onnx_path)
    tf_rep = backend.prepare(onnx_model)
    tf_rep.export_graph(tf_dir)


if __name__ == '__main__':
    onnx2tf(onnx_path='./model.onnx', tf_dir='./model/001/')

TensorFlow Serving模型部署

由于docker只能启动一个TensorFlow Serving服务，当有多个模型需要部署时需要在模型文件中添加一个配置文件，单个模型部署也建议添加配置文件，方便以后添加模型。
要求模型文件（model1，model2）与配置文件（model.config）在同一个目录下：

models/
|-- model1
|   `-- 001
|       |-- assets
|       |-- saved_model.pb
|       `-- variables
|           |-- variables.data-00000-of-00001
|           `-- variables.index
|-- model2
|   `-- 001
|       |-- assets
|       |-- saved_model.pb
|       `-- variables
|           |-- variables.data-00000-of-00001
|           `-- variables.index
|-- model.config

配置文件model.config格式

model_config_list: {
  config: {
    name: "model1",
    base_path: "/models/model1/",
    model_platform: "tensorflow"
  },
  config: {
    name: "model2",
    base_path: "/models/model2/",
    model_platform: "tensorflow"
  }
}

然后运行下面命令即可启动TensorFlow Serving服务

sudo docker run -p 8501:8501 --name tf_serving --mount type=bind,source=models目录路径,target=/models/ -t tensorflow/serving --model_config_file=/models/model.config &

为了确定模型是否已经部署成功，需要在本地调用TensorFlow Serving进行测试

import requests
url = "http://IP地址:8501/v1/models/model1:predict"
img = 经过预处理的模型输入
predict_request = {'instances': img.tolist()}
response = requests.post(url=url, json=predict_request)
predictions = response.json()['predictions']  # 未经后处理的模型输出，类型为list

调用TensorFlow Serving提供接口

由于TensorFlow Serving不能对数据进行前处理和后处理，实际应用中无法直接调用，所以需要再写一个服务来进行数据的预处理和后处理，这里使用python的flask：1、接受一张图片，2、对其进行预处理，3、调用TensorFlow Serving对处理后的图片目标检测，4、对目标检测的输出进行后处理，5、返回最终的结果。

服务端

import requests
from PIL import Image
import numpy as np
from flask import Flask, request, jsonify

app = Flask(__name__)


@app.route('/object_detection', methods=['POST'])
def img2bbox():
    file = request.files['file']  # 接受一张图片文件
    img_rgb = Image.open(file).convert('RGB')  # 读取图片文件
    img_arr = np.array(img_rgb)  # 转成<numpy.ndarray>类型
    img = preprocess(img_arr)  # 图片预处理
    predict_request = {'instances': img.tolist()}
	response = requests.post(url="http://IP地址:8501/v1/models/model1:predict", json=predict_request)
	predictions = response.json()['predictions']  # 未经后处理的模型输出，类型为list
	predictions = np.array(predictions)  # 转成<numpy.ndarray>类型
	res = postprocess(predictions)  # 后处理
    return res


if __name__ == "__main__":
    host = '0.0.0.0'
    port = 1234
    app.run(host=host, port=port, debug=True, use_reloader=False)

客户端

import requests

with open(img_path, 'rb') as f:
    file = f.read()

res = requests.post(url=url, files={'file': file})
print(response_time)
result = res.json()