使用 tf.lookup.StaticHashTable 导出 TF2.2.0 模型进行服务时出错

Question

StaticHashTable我在 tf.keras 模型的输出层之后的一个 Lambda 层中使用as 。实际上非常简单：我有一个文本分类模型，并且添加了一个简单的 lambda 层，该层接受model.outputmodel_id 并将其转换为更通用的标签。我可以使用 model.save(... as H5 format..) 保存此版本的模型，没有任何问题，并且可以加载它并使用它，没有任何问题。

问题是，当我尝试导出 TF-Serving 的 TF2.2.0 模型时，我找不到如何导出它。这是我可以用 TF1.X 或用TF2.X + tf.compat.v1.disable_eager_execution()

tf.compat.v1.disable_eager_execution()
version = 1
name = 'tmp_model'
export_path = f'/opt/tf_serving/{name}/{version}'
builder = saved_model_builder.SavedModelBuilder(export_path)

model_signature = tf.compat.v1.saved_model.predict_signature_def(
    inputs={
        'input': model.input
    }, 
    outputs={
        'output': model.output
    }
)

with tf.compat.v1.keras.backend.get_session() as sess:
    builder.add_meta_graph_and_variables(
        sess=sess,
        tags=[tf.compat.v1.saved_model.tag_constants.SERVING],
        signature_def_map={
            'predict': model_signature
        },
        # For initializing Hashtables
        main_op=tf.compat.v1.tables_initializer()
    )
    builder.save()

这将以 TF1.X 格式保存我的模型以供服务，我可以毫无问题地使用它。事情是，我正在使用 LSTM 层，并且我想在 GPU 上使用我的模型。根据文档，如果我禁用 eager 模式，则无法将 GPU 版本的 LSTM 与 TF2.2 一起使用。如果不经过上述代码，我无法保存我的模型以服务于 TF2.2 标准和 StaticHashTables。

以下是我尝试导出在最后一层使用 StaticHashTables 的 TF2.2 模型的方法；并给出如下错误：

class MyModule(tf.Module):

    def __init__(self, model):
        super(MyModule, self).__init__()
        self.model = model
    
    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 16), dtype=tf.int32, name='input')])
    def predict(self, input):
        result = self.model(input)
        return {"output": result}

version = 1
name = 'tmp_model'
export_path = f'/opt/tf_serving/{name}/{version}'

module = MyModule(model)
tf.saved_model.save(module, export_path, signatures={"predict": module.predict.get_concrete_function()})

错误：

AssertionError: Tried to export a function which references untracked object Tensor("2907:0", shape=(), dtype=resource).
TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.

对于导出使用StaticHashTables最终 Lambda 层进行 TensorFlow Serving 的 TF2.2 模型，有什么建议吗？或者我是否遗漏了任何内容？

更多信息在这里：https://github.com/tensorflow/serving/issues/1719

谢谢！

Answer 1

我遇到了同样的问题，我找到了答案，通过查找转换创建自定义层，然后将该层添加到我的模型中。其他人把答案放在 stackoverflow 上，但我找不到了，所以我会把它给你。原因是来自外部的变量和其他元素必须是可跟踪的，我没有找到其他方法使它们可跟踪，而是创建一个自定义图层，因为它们是可跟踪的，并且在导出时不需要添加额外的资源。

\n

这是代码：

\n

这是特定于在模型之前进行转换的自定义层（包括作为从静态表查找的标记生成器，然后是填充：

\n

class VocabLookup(tf.keras.layers.Layer):\n    def __init__(self, word_index, **kwargs):\n        self.word_index = word_index\n        self.vocab = list(word_index.keys())\n        self.indices = tf.convert_to_tensor(list(word_index.values()), dtype=tf.int64)\n        vocab_initializer = tf.lookup.KeyValueTensorInitializer(self.vocab, self.indices)\n        self.table = tf.lookup.StaticHashTable(vocab_initializer, default_value=1)\n        super(VocabLookup, self).__init__(**kwargs)\n\n    def build(self, input_shape):\n        self.built = True\n\n    def sentences_transform(self,tx):\n        x = tf.strings.lower(tx)\n        x = tf.strings.regex_replace(x,"[,.:;]", " ")\n        x = tf.strings.regex_replace(x,"\xc3\xa1", "a")\n        x = tf.strings.regex_replace(x,"\xc3\xa9", "e")\n        x = tf.strings.regex_replace(x,"\xc3\xad", "i")\n        x = tf.strings.regex_replace(x,"\xc3\xb3", "i")\n        x = tf.strings.regex_replace(x,"\xc3\xba", "u")\n        x = tf.strings.regex_replace(x,"\xc3\xbc", "u")\n        x = tf.strings.regex_replace(x,"\xc3\x81", "a")\n        x = tf.strings.regex_replace(x,"\xc3\x89", "e")\n        x = tf.strings.regex_replace(x,"\xc3\x8d", "i")\n        x = tf.strings.regex_replace(x,"\xc3\x93", "o")\n        x = tf.strings.regex_replace(x,"\xc3\x9a", "u")\n        x = tf.strings.regex_replace(x,"\xc3\x9c", "u")\n        x = tf.strings.regex_replace(x,"\xc3\x9c", "u")\n        x = tf.strings.regex_replace(x,"[?\xc2\xbf\xc2\xa1!@#$-_\\?+\xc2\xbf{}*/]", "")\n        x = tf.strings.regex_replace(x," +", " ")\n        x = tf.strings.strip(x)\n        x = tf.strings.split(x)\n        x = self.table.lookup(x)\n        x_as_vector = tf.reshape(x, [-1])\n        zero_padding = tf.zeros([191] - tf.shape(x_as_vector), dtype=x.dtype)\n        x = tf.concat([x_as_vector, zero_padding], 0)\n        return x\n        \n\n    def call(self, inputs):\n        x = tf.map_fn(lambda tx: self.sentences_transform(tx), elems = inputs,dtype=tf.int64)\n        return x\n\n    def get_config(self):\n        return {'word_index': self.word_index}\n

\n

在我的例子中，我创建了一个层来从分词器接收 word_index 作为输入。然后，您可以在模型中使用如下所示的图层：

\n

with open(<tokenizer_path>) as f:\n    data = json.load(f)\n    tokenizer = tokenizer_from_json(data)\n\nmoderator = load_model(<final model path ('.h5')>)\nword_index = tokenizer.word_index\ntext_bytes = tf.keras.Input(shape=(), name='image_bytes', dtype=tf.string)\nx = VocabLookup(word_index)(text_bytes)\noutput = moderator(x)\nmodel = tf.keras.models.Model(text_bytes, output)\n

\n

如果您进行总结，您将得到如下内容：

\n

model.summary()\nModel: "functional_57"\n_________________________________________________________________\nLayer (type)                 Output Shape              Param #   \n=================================================================\nimage_bytes (InputLayer)     [(None,)]                 0         \n_________________________________________________________________\nvocab_lookup_60 (VocabLookup (None, None)              0         \n_________________________________________________________________\nsequential_1 (Sequential)    (None, 1)                 1354369   \n=================================================================\nTotal params: 1,354,369\nTrainable params: 1,354,369\nNon-trainable params: 0\n

\n

通过此步骤，您终于可以另存为 TF2 服务模型

\n

save_path = <your_serving_model_path>\ntf.saved_model.save(model,  save_path)\n

\n