Python调用TensorFlow Serving部署的模型提供接口
安装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()