TensorFlow Serving模型部署

TensorFlow Serving可以直接读取SavedModel格式的模型进行部署(导出模型到SavedModel文件的方法见 前文 )。使用以下命令即可:

  1. tensorflow_model_server \
  2. --rest_api_port=端口号(如8501 \
  3. --model_name=模型名 \
  4. --model_base_path="SavedModel格式模型的文件夹绝对地址(不含版本号)"

注解

TensorFlow Serving支持热更新模型,其典型的模型文件夹结构如下:

  1. /saved_model_files
  2. /1 # 版本号为1的模型文件
  3. /assets
  4. /variables
  5. saved_model.pb
  6. ...
  7. /N # 版本号为N的模型文件
  8. /assets
  9. /variables
  10. saved_model.pb

上面1~N的子文件夹代表不同版本号的模型。当指定 —model_base_path 时,只需要指定根目录的 绝对地址 (不是相对地址)即可。例如,如果上述文件夹结构存放在 home/snowkylin 文件夹内,则 —model_base_path 应当设置为 home/snowkylin/saved_model_files (不附带模型版本号)。TensorFlow Serving会自动选择版本号最大的模型进行载入。

Keras Sequential模式模型的部署

由于Sequential模式的输入和输出都很固定,因此这种类型的模型很容易部署,无需其他额外操作。例如,要将 前文使用SavedModel导出的MNIST手写体识别模型 (使用Keras Sequential模式建立)以 MLP 的模型名在 8501 端口进行部署,可以直接使用以下命令:

  1. tensorflow_model_server \
  2. --rest_api_port=8501 \
  3. --model_name=MLP \
  4. --model_base_path="/home/.../.../saved" # 文件夹绝对地址根据自身情况填写,无需加入版本号

然后就可以按照 后文的介绍 ,使用gRPC或者RESTful API在客户端调用模型了。

自定义Keras模型的部署

使用继承 tf.keras.Model 类建立的自定义Keras模型的自由度相对更高。因此当使用TensorFlow Serving部署模型时,对导出的SavedModel文件也有更多的要求:

  • 需要导出到SavedModel格式的方法(比如 call )不仅需要使用 @tf.function 修饰,还要在修饰时指定 input_signature 参数,以显式说明输入的形状。该参数传入一个由 tf.TensorSpec 组成的列表,指定每个输入张量的形状和类型。例如,对于MNIST手写体数字识别,我们的输入是一个 [None, 28, 28, 1] 的四维张量( None 表示第一维即Batch Size的大小不固定),此时我们可以将模型的 call 方法做以下修饰:
  1. class MLP(tf.keras.Model):
  2. ...
  3.  
  4. @tf.function(input_signature=[tf.TensorSpec([None, 28, 28, 1], tf.float32)])
  5. def call(self, inputs):
  6. ...
  • 在将模型使用 tf.saved_model.save 导出时,需要通过 signature 参数提供待导出的函数的签名(Signature)。简单说来,由于自定义的模型类里可能有多个方法都需要导出,因此,需要告诉TensorFlow Serving每个方法在被客户端调用时分别叫做什么名字。例如,如果我们希望客户端在调用模型时使用 call 这一签名来调用 model.call 方法时,我们可以在导出时传入 signature 参数,以 dict 的键值对形式告知导出的方法对应的签名,代码如下:
  1. model = MLP()
  2. ...
  3. tf.saved_model.save(model, "saved_with_signature/1", signatures={"call": model.call})

以上两步均完成后,即可使用以下命令部署:

  1. tensorflow_model_server \
  2. --rest_api_port=8501 \
  3. --model_name=MLP \
  4. --model_base_path="/home/.../.../saved_with_signature" # 修改为自己模型的绝对地址