TensorFlow Serving 可以快速部署 Tensorflow 模型，上线 gRPC 或 REST API。
官方推荐 Docker 部署，也给了训练到部署的完整教程：Servers: TFX for TensorFlow Serving。本文只是遵照教程进行的练习，有助于了解 TensorFlow 训练到部署的整个过程。
准备环境
准备好 TensorFlow 环境，导入依赖：
1import sys 2 3# Confirm that we're using Python 3 4assert sys.version_info.major == 3, 'Oops, not running Python 3. Use Runtime > Change runtime type'
1import tensorflow as tf 2from tensorflow import keras 3 4# Helper libraries 5import numpy as np 6import matplotlib.pyplot as plt 7import os 8import subprocess 9 10print(f'TensorFlow version: {tf.__version__}') 11print(f'TensorFlow GPU support: {tf.test.is_built_with_gpu_support()}') 12 13physical_gpus = tf.config.list_physical_devices('GPU') 14print(physical_gpus) 15for gpu in physical_gpus: 16 # memory growth must be set before GPUs have been initialized 17 tf.config.experimental.set_memory_growth(gpu, True) 18logical_gpus = tf.config.experimental.list_logical_devices('GPU') 19print(len(physical_gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
1TensorFlow version: 2.4.1 2TensorFlow GPU support: True 3[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 41 Physical GPUs, 1 Logical GPUs
创建模型
载入 Fashion MNIST 数据集：
1fashion_mnist = keras.datasets.fashion_mnist 2(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data() 3 4# scale the values to 0.0 to 1.0 5train_images = train_images / 255.0 6test_images = test_images / 255.0 7 8# reshape for feeding into the model 9train_images = train_images.reshape(train_images.shape[0], 28, 28, 1) 10test_images = test_images.reshape(test_images.shape[0], 28, 28, 1) 11 12class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 13 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'] 14 15print('\ntrain_images.shape: {}, of {}'.format(train_images.shape, train_images.dtype)) 16print('test_images.shape: {}, of {}'.format(test_images.shape, test_images.dtype))
1train_images.shape: (60000, 28, 28, 1), of float64 2test_images.shape: (10000, 28, 28, 1), of float64
用最简单的 CNN 训练模型，
1model = keras.Sequential([ 2 keras.layers.Conv2D(input_shape=(28,28,1), filters=8, kernel_size=3, 3 strides=2, activation='relu', name='Conv1'), 4 keras.layers.Flatten(), 5 keras.layers.Dense(10, name='Dense') 6]) 7model.summary() 8 9testing = False 10epochs = 5 11 12model.compile(optimizer='adam', 13 loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 14 metrics=[keras.metrics.SparseCategoricalAccuracy()]) 15model.fit(train_images, train_labels, epochs=epochs) 16 17test_loss, test_acc = model.evaluate(test_images, test_labels) 18print('\nTest accuracy: {}'.format(test_acc))
1Model: "sequential" 2_________________________________________________________________ 3Layer (type) Output Shape Param # 4================================================================= 5Conv1 (Conv2D) (None, 13, 13, 8) 80 6_________________________________________________________________ 7flatten (Flatten) (None, 1352) 0 8_________________________________________________________________ 9Dense (Dense) (None, 10) 13530 10================================================================= 11Total params: 13,610 12Trainable params: 13,610 13Non-trainable params: 0 14_________________________________________________________________ 15Epoch 1/5 161875/1875 [==============================] - 3s 722us/step - loss: 0.7387 - sparse_categorical_accuracy: 0.7449 17Epoch 2/5 181875/1875 [==============================] - 1s 793us/step - loss: 0.4561 - sparse_categorical_accuracy: 0.8408 19Epoch 3/5 201875/1875 [==============================] - 1s 720us/step - loss: 0.4097 - sparse_categorical_accuracy: 0.8566 21Epoch 4/5 221875/1875 [==============================] - 1s 718us/step - loss: 0.3899 - sparse_categorical_accuracy: 0.8636 23Epoch 5/5 241875/1875 [==============================] - 1s 719us/step - loss: 0.3673 - sparse_categorical_accuracy: 0.8701 25313/313 [==============================] - 0s 782us/step - loss: 0.3937 - sparse_categorical_accuracy: 0.8630 26 27Test accuracy: 0.8629999756813049
保存模型
将模型保存成 SavedModel 格式，路径里加上版本号，以便 TensorFlow Serving 时可选择模型版本。
1# Fetch the Keras session and save the model 2# The signature definition is defined by the input and output tensors, 3# and stored with the default serving key 4import tempfile 5 6MODEL_DIR = os.path.join(tempfile.gettempdir(), 'tfx') 7version = 1 8export_path = os.path.join(MODEL_DIR, str(version)) 9print('export_path = {}\n'.format(export_path)) 10 11tf.keras.models.save_model( 12 model, 13 export_path, 14 overwrite=True, 15 include_optimizer=True, 16 save_format=None, 17 signatures=None, 18 options=None 19) 20 21print('\nSaved model:') 22!ls -l {export_path}
1export_path = /tmp/tfx/1 2 3INFO:tensorflow:Assets written to: /tmp/tfx/1/assets 4 5Saved model: 6total 88 7drwxr-xr-x 2 john john 4096 Apr 13 15:10 assets 8-rw-rw-r-- 1 john john 78169 Apr 13 15:12 saved_model.pb 9drwxr-xr-x 2 john john 4096 Apr 13 15:12 variables
查看模型
使用 saved_model_cli 工具查看模型的 MetaGraphDefs (the models) 和 SignatureDefs (the methods you can call)，了解信息。
1!saved_model_cli show --dir '/tmp/tfx/1' --all
12021-04-13 15:12:29.433576: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0 2 3MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs: 4 5signature_def['__saved_model_init_op']: 6 The given SavedModel SignatureDef contains the following input(s): 7 The given SavedModel SignatureDef contains the following output(s): 8 outputs['__saved_model_init_op'] tensor_info: 9 dtype: DT_INVALID 10 shape: unknown_rank 11 name: NoOp 12 Method name is: 13 14signature_def['serving_default']: 15 The given SavedModel SignatureDef contains the following input(s): 16 inputs['Conv1_input'] tensor_info: 17 dtype: DT_FLOAT 18 shape: (-1, 28, 28, 1) 19 name: serving_default_Conv1_input:0 20 The given SavedModel SignatureDef contains the following output(s): 21 outputs['Dense'] tensor_info: 22 dtype: DT_FLOAT 23 shape: (-1, 10) 24 name: StatefulPartitionedCall:0 25 Method name is: tensorflow/serving/predict 26 27Defined Functions: 28 Function Name: '__call__' 29 Option #1 30 Callable with: 31 Argument #1 32 Conv1_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='Conv1_input') 33 Argument #2 34 DType: bool 35 Value: False 36 Argument #3 37 DType: NoneType 38 Value: None 39 Option #2 40 Callable with: 41 Argument #1 42 inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs') 43 Argument #2 44 DType: bool 45 Value: False 46 Argument #3 47 DType: NoneType 48 Value: None 49 Option #3 50 Callable with: 51 Argument #1 52 inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs') 53 Argument #2 54 DType: bool 55 Value: True 56 Argument #3 57 DType: NoneType 58 Value: None 59 Option #4 60 Callable with: 61 Argument #1 62 Conv1_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='Conv1_input') 63 Argument #2 64 DType: bool 65 Value: True 66 Argument #3 67 DType: NoneType 68 Value: None 69 ...
部署模型
安装 Serving
1echo "deb [arch=amd64] http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-model-server tensorflow-model-server-universal" | sudo tee /etc/apt/sources.list.d/tensorflow-serving.list && \ 2curl https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | sudo apt-key add - 3 4sudo apt update 5sudo apt install tensorflow-model-server
开启 Serving
开启 TensorFlow Serving ，提供 REST API ：
- rest_api_port: REST 请求端口。
- model_name: REST 请求 URL ，自定义的名称。
- model_base_path: 模型所在目录。
1nohup tensorflow_model_server \ 2 --rest_api_port=8501 \ 3 --model_name=fashion_model \ 4 --model_base_path="/tmp/tfx" >server.log 2>&1 &
1$ tail server.log 2To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags. 32021-04-13 15:12:10.706648: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:206] Restoring SavedModel bundle. 42021-04-13 15:12:10.726722: I external/org_tensorflow/tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 2599990000 Hz 52021-04-13 15:12:10.756506: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:190] Running initialization op on SavedModel bundle at path: /tmp/tfx/1 62021-04-13 15:12:10.759935: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:277] SavedModel load for tags { serve }; Status: success: OK. Took 110653 microseconds. 72021-04-13 15:12:10.760277: I tensorflow_serving/servables/tensorflow/saved_model_warmup_util.cc:59] No warmup data file found at /tmp/tfx/1/assets.extra/tf_serving_warmup_requests 82021-04-13 15:12:10.760486: I tensorflow_serving/core/loader_harness.cc:87] Successfully loaded servable version {name: fashion_model version: 1} 92021-04-13 15:12:10.763938: I tensorflow_serving/model_servers/server.cc:371] Running gRPC ModelServer at 0.0.0.0:8500 ... 10[evhttp_server.cc : 238] NET_LOG: Entering the event loop ... 112021-04-13 15:12:10.765308: I tensorflow_serving/model_servers/server.cc:391] Exporting HTTP/REST API at:localhost:8501 ...
访问服务
随机显示一张测试图：
1def show(idx, title): 2 plt.figure() 3 plt.imshow(test_images[idx].reshape(28,28)) 4 plt.axis('off') 5 plt.title('\n\n{}'.format(title), fontdict={'size': 16}) 6 7import random 8rando = random.randint(0,len(test_images)-1) 9show(rando, 'An Example Image: {}'.format(class_names[test_labels[rando]]))
创建 JSON 对象，给到三张要预测的图：
1import json 2data = json.dumps({"signature_name": "serving_default", "instances": test_images[0:3].tolist()}) 3print('Data: {} ... {}'.format(data[:50], data[len(data)-52:]))
1Data: {"signature_name": "serving_default", "instances": ... [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]]]]}
REST 请求
最新模型版本进行预测：
1!pip install -q requests 2 3import requests 4headers = {"content-type": "application/json"} 5json_response = requests.post('http://localhost:8501/v1/models/fashion_model:predict', data=data, headers=headers) 6predictions = json.loads(json_response.text)['predictions'] 7 8show(0, 'The model thought this was a {} (class {}), and it was actually a {} (class {})'.format( 9 class_names[np.argmax(predictions[0])], np.argmax(predictions[0]), class_names[test_labels[0]], test_labels[0]))
指定模型版本进行预测：
1headers = {"content-type": "application/json"} 2json_response = requests.post('http://localhost:8501/v1/models/fashion_model/versions/1:predict', data=data, headers=headers) 3predictions = json.loads(json_response.text)['predictions'] 4 5for i in range(0,3): 6 show(i, 'The model thought this was a {} (class {}), and it was actually a {} (class {})'.format( 7 class_names[np.argmax(predictions[i])], np.argmax(predictions[i]), class_names[test_labels[i]], test_labels[i]))
GoCoding 个人实践的经验分享，可关注公众号！