如何使用 tensorflow 服务使 tensorflow hub 嵌入可服务？答案

【问题标题】：How to make the tensorflow hub embeddings servable using tensorflow serving?如何使用 tensorflow 服务使 tensorflow hub 嵌入可服务？
【发布时间】：2020-12-21 22:04:41
【问题描述】：

我正在尝试使用来自 tensorflow hub 的嵌入模块作为可服务。我是张量流的新手。目前，我使用Universal Sentence Encoder 嵌入作为查找将句子转换为嵌入，然后使用这些嵌入来查找与另一个句子的相似性。

我当前将句子转换为嵌入的代码是：

with tf.Session() as session:
  session.run([tf.global_variables_initializer(), tf.tables_initializer()])
  sen_embeddings = session.run(self.embed(prepared_text))

Prepared_text 是一个句子列表。我如何采用这个模型并使其成为可服务的？

【问题讨论】：

标签： tensorflow tensorflow-serving word-embedding

【解决方案1】：

现在您可能需要手动执行此操作。这是我的解决方案，类似于之前的答案，但更通用 - 展示如何在不猜测输入参数的情况下使用任何其他模块，以及通过验证和使用进行扩展：

import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.saved_model import simple_save

export_dir = "/tmp/tfserving/universal_encoder/00000001"
with tf.Session(graph=tf.Graph()) as sess:
    module = hub.Module("https://tfhub.dev/google/universal-sentence-encoder/2") 
    input_params = module.get_input_info_dict()
    # take a look at what tensor does the model accepts - 'text' is input tensor name

    text_input = tf.placeholder(name='text', dtype=input_params['text'].dtype, 
        shape=input_params['text'].get_shape())
    sess.run([tf.global_variables_initializer(), tf.tables_initializer()])

    embeddings = module(text_input)

    simple_save(sess,
        export_dir,
        inputs={'text': text_input},
        outputs={'embeddings': embeddings},
        legacy_init_op=tf.tables_initializer())

感谢module.get_input_info_dict()，您知道需要将哪些张量名称传递给模型——您在simple_save 方法中将此名称用作inputs={} 的键。

请记住，要为模型提供服务，它需要位于以 version 结尾的目录路径中，这就是为什么 '00000001' 是 saved_model.pb 所在的最后一个路径。

导出模块后，查看模型是否正确导出以供服务的最快方法是使用saved_model_cli API：

saved_model_cli run --dir /tmp/tfserving/universal_encoder/00000001 --tag_set serve --signature_def serving_default --input_exprs 'text=["what this is"]'

从 docker 提供模型：

docker pull tensorflow/serving  
docker run -p 8501:8501 -v /tmp/tfserving/universal_encoder:/models/universal_encoder -e MODEL_NAME=universal_encoder -t tensorflow/serving

【讨论】：

【解决方案2】：

目前，Tensorflow Serving 无法直接使用集线器模块。您必须将模块加载到一个空图表中，然后使用SavedModelBuilder 将其导出。例如：

import tensorflow as tf
import tensorflow_hub as hub

with tf.Graph().as_default():
  module = hub.Module("http://tfhub.dev/google/universal-sentence-encoder/2")
  text = tf.placeholder(tf.string, [None])
  embedding = module(text)

  init_op = tf.group([tf.global_variables_initializer(), tf.tables_initializer()])
  with tf.Session() as session:
    session.run(init_op)
    tf.saved_model.simple_save(
        session,
        "/tmp/serving_saved_model",
        inputs = {"text": text},
        outputs = {"embedding": embedding},
        legacy_init_op = tf.tables_initializer()        
    )

这将以所需的格式导出您的模型（到文件夹/tmp/serving_saved_model）以进行投放。之后，您可以按照此处文档中的说明进行操作：https://www.tensorflow.org/serving/serving_basic

【讨论】：

【解决方案3】：

请注意，其他答案适用于 TensorFlow 1。TensorFlow 2 的大多数 TF Hub 模型已经与 TF Serving 兼容。例如，部署 USE-Large 模型：

通过tensorflow_hub 库或仅通过https://tfhub.dev/google/universal-sentence-encoder-large/5 下载模型
将内容放入代表型号名称和版本的文件夹中，例如models/use-large/5
运行 TF Serving 应用程序，例如通过 Docker：

docker run -t --rm -p 8501:8501 \
   -v "$PATH_TO_YOUR_WORKSPACE/models:/models" \
   -e MODEL_NAME="use-large" \
   tensorflow/serving

该模型将在localhost:8501/v1/models/use-large：

curl -d '{"instances": ["Hey!"]}' \
    -X POST http://localhost:8501/v1/models/use-large:predict

【讨论】：