這篇文章主要介紹TF Serving
的使用,主要用於在生產環境中部署TensorFlow
模型;
分成四個部分進行介紹:
TF Serving
概述:介紹其基本概念、作用、原理;- 模型准備:通過自定義一個簡單的神經網絡,包含數據准備、模型訓練,並導出所需模型;
- 服務器端:編寫服務器端相關程序;
- 客戶端:在客戶端如何使用服務器端提供的模型服務;
- 多模型部署:
以上五個部分對應1-5節內容;
注意:這篇文章只是提供了基本功能的實現,很多原理,還不理解;先這樣,后續再更新;
項目地址:tfserving
1. TF Serving概述
TF Serving
是一個用於機器學習模型部署的高性能開源庫,可以將訓練完成的模型直接上線並提供服務;
一個重要的特性是:支持熱更新與模型版本自動管理,這意味着一旦部署成功,不需要再為線上服務擔心,只需要關心線下的模型訓練即可;
核心模塊如下圖所示:
Source
用於管理本地的模型文件系統;Source
會對需要加載的模型創建一個Loader
,Loder
中會包含要加載模型的全部信息;Source
通知Manager
有新的模型需要進行加載;Manager
通過版本管理策略Version Policy
來確定哪些模型需要被下架,哪些模型需要被加載;Manager
在確認需要加載的模型是否符合加載策略,便通知Loader
來加載最新的模型;Client
向服務器端請求模型結果,可以指定版本;
2. 模型准備
這一部分主要目的是為了得到訓練完成的模型文件,類型為.pb
格式;
這里提供一個簡單的程序,演示從數據准備、模型訓練、模型保存的過程中;
其中,重點是最后一步模型保存,需要保存成pb
格式,如果你這邊是其他格式,需要將其轉換成pb
;
另外,在TensorFlow 2.0+
版本中,model.save()
默認的存儲格式就是pb
形式;
訓練完成后,模型文件如下:
- 其中,
clothing
表示該模型用於識別衣服類型;(可自定義名稱) 1
表示clothing
模型的版本號,目的是為了后續更新模型;
完成程序model_training.py
如下:
# -*- coding: utf-8 -*-
# =================================================
"""
2. Model Prepare:
Train a neural network model to classify image of clothing, like sneakers and shirts.
Finally, output a '.pb' model
"""
import tensorflow as tf
# =========================
# ===== Load dataset ======
# =========================
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
print("train_images: {}, train_labels: {}".format(train_images.shape, train_labels.shape))
print("test_images: {}, test_labels: {}".format(test_images.shape, test_labels.shape))
# =========================
# ====== Preprocess =======
# =========================
train_images = train_images / 255.0
test_images = test_images / 255.0
# =========================
# ==== Build the model ====
# =========================
model = tf.keras.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
# =========================
# ==== Train the model ====
# =========================
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=10)
# =========================
# ==== Save the model ====
# =========================
model.save('./saved_model/clothing/1/', save_format='tf') # save_format: Defaults to 'tf' in TF 2.X, and 'h5' in TF 1.X.
3. 服務器端
服務器端的目的是基於上面的模型文件,提供模型服務,方便通過客戶端進行調用,實現模型推理;
這里提供的方式是通過利用docker
進行部署;
3.1 docker安裝
首先,確定docker
是否安裝:
sudo docker --version
得到:
Docker version 19.03.6, build 369ce74a3c
如果沒有安裝docker
,請先安裝;
3.2 拉取docker-serving鏡像
sudo docker pull tensorflow/serving:2.3.0
完成后,通過下面命令查看:
sudo docker images
3.3 啟動TF Serving服務
創建一個運行在后台服務的容器
sudo docker run -p 8500:8500 -p 8501:8501 --mount type=bind,source=/***/saved_model/clothing,target=/models/clothing -e MODEL_NAME=clothing -t tensorflow/serving:2.3.0
-p 8500:8500 -p 8501:8501
表示將本地的端口映射到容器的端口,其中8500
表示通信方式使用gPRC
,8501
使用Rest API
進行通信,二選一,這里選擇8500
;--mount
表示將本地目錄掛載到容器內部;source
表示模型文件本地存儲的絕對路徑,到clothing
這一級目錄;target
表示模型文件被映射到容器內部的路徑;MODEL_NAME
表示模型名稱,用於在發送POST請求時唯一標識符;--name
表示給容器賦予一個名稱,后續對容器的關閉、刪除、重啟,都可以使用該名稱;-t
表示Docker為要創建的容器分配一個偽tty終端;
啟動成功后,可以通過下面命令查看容器狀態:
sudo docker ps -a
4. 客戶端
編寫客戶端程序,完成對下面這張圖片的預測;
4.1 查看網絡輸入輸出節點
使用下面的命令查看網絡模型的輸入、輸出節點名稱,也可以在定義網絡時指明節點名稱:
saved_model_cli show --dir='saved_model/clothing/1/' --all
4.2 客戶端推理
完整程序如下:
# -*- coding: utf-8 -*-
# =================================================
import numpy as np
from PIL import Image
from deploy import PredictModelGrpc
"""
4. Client
"""
# =======================
# ===== Load image ======
# =======================
image_path = './test_images/img_1.png'
img = Image.open(image_path)
img = np.array(img) / 255.0
# =======================
# ===== Load image ======
# =======================
model = PredictModelGrpc(model_name='clothing', input_name='flatten_input', output_name='dense_1')
res = model.inference(img)
print(res)
其中,PredictModelGrpc
定義在deploy.py
中,如下:
# -*- coding: utf-8 -*-
# =================================================
import tensorflow as tf
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import grpc
import numpy as np
class PredictModelGrpc(object):
def __init__(self, model_name, input_name, output_name, socket='0.0.0.0:8500'):
self.socket = socket
self.model_name = model_name
self.input_name = input_name
self.output_name = output_name
self.request, self.stub = self.__get_request()
def __get_request(self):
channel = grpc.insecure_channel(self.socket, options=[('grpc.max_send_message_length', 1024 * 1024 * 1024),
('grpc.max_receive_message_length',
1024 * 1024 * 1024)]) # 可設置大小
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
request = predict_pb2.PredictRequest()
request.model_spec.name = self.model_name
request.model_spec.signature_name = "serving_default"
return request, stub
def inference(self, frames):
self.request.inputs[self.input_name].CopyFrom(tf.make_tensor_proto(frames, dtype=tf.float32)) # images is input of model
result = self.stub.Predict.future(self.request, 10.0)
res = tf.make_ndarray(result.result().outputs[self.output_name])[0]
return res
if __name__ == '__main__':
model = PredictModelGrpc(model_name='03_win_day')
import time
for i in range(10):
s = time.time()
data = np.random.uniform(0, 1, (1, 12, 3, 3, 11))
result = model.inference(data)
print(result.shape)
e = time.time()
print(e-s)
4. 多模型部署
在實際部署的時候,可能包含有多個模型,不同之處在於兩點:
- 啟動
TF Serving
時,需要指定一個config
文件,用於指定多個模型的路徑; - 客戶端使用時,調用不同的模型,需要指定具體的
model_name
;
4.1 多個模型
如下圖所示,假設存在兩個模型clothing
和clothing_2
:
4.2 多模型配置文件
在saved_model_multi
目錄下,創建models.config
文件,用於啟動TF Serving
時使用;
model_config_list:{
config:{
name: "clothing",
base_path: "/models/saved_model_multi/clothing",
model_platform: "tensorflow",
model_version_policy:{
all: {}
}
},
config:{
name: "clothing_2",
base_path: "/models/saved_model_multi/clothing_2",
model_platform: "tensorflow",
}
}
- 其中,
model_config_list
包含多個config
,用於配置每一個模型; name
表示模型唯一標識符;base_path
表示當前模型在容器內部的路徑,與docker run
掛載到容器內目錄保持一致;
4.3 啟動TF Serving服務
sudo docker run -p 8500:8500 --mount type=bind,source=/data1/chenz/workspace/tfserving/saved_model_multi,target=/models/saved_model_multi --name=tfserving_test_multi -t tensorflow/serving:2.3.0 --model_config=/models/saved_model/models.config &
-p 8500:8500
將本地端口映射到容器端口;--mount
掛載命令,將本地目錄掛載到容器內部;type=bind
掛載類型;source
本地目錄;target
容器內部目錄;--name
容器別名,方便對容器進行操作:停止、刪除、重啟;-t
使用的鏡像;--model_config
多模型部署時的配置文件;
4.4 客戶端使用
在客戶端使用時,只需要根據自定義的條件,給PredictModelGrpc
傳遞不同的參數即可:
model = PredictModelGrpc(model_name='clothing', input_name='flatten_input', output_name='dense_1')
或者:
model = PredictModelGrpc(model_name='clothing_2', input_name='flatten_input', output_name='dense_1')