时间戳

首页 > 解决方案 > 使用 tf.lookup.StaticHashTable 导出 TF2.2.0 模型进行服务时出错

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

问题描述

我在StaticHashTable我的 tf.keras 模型的输出层之后的一个 Lambda 层中使用 as。实际上很简单：我有一个文本分类模型，我正在添加一个简单的 lambda 层，该层model.output将 model_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 模式，我将无法将 LSTM 的 GPU 版本与 TF2.2 一起使用。如果不通过上述代码，我无法保存我的模型来服务 wrt 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

谢谢！

标签： tensorflowkerastensorflow2.0tf.kerastensorflow-serving

解决方案

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

这是代码：

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

class VocabLookup(tf.keras.layers.Layer):
    def __init__(self, word_index, **kwargs):
        self.word_index = word_index
        self.vocab = list(word_index.keys())
        self.indices = tf.convert_to_tensor(list(word_index.values()), dtype=tf.int64)
        vocab_initializer = tf.lookup.KeyValueTensorInitializer(self.vocab, self.indices)
        self.table = tf.lookup.StaticHashTable(vocab_initializer, default_value=1)
        super(VocabLookup, self).__init__(**kwargs)

    def build(self, input_shape):
        self.built = True

    def sentences_transform(self,tx):
        x = tf.strings.lower(tx)
        x = tf.strings.regex_replace(x,"[,.:;]", " ")
        x = tf.strings.regex_replace(x,"á", "a")
        x = tf.strings.regex_replace(x,"é", "e")
        x = tf.strings.regex_replace(x,"í", "i")
        x = tf.strings.regex_replace(x,"ó", "i")
        x = tf.strings.regex_replace(x,"ú", "u")
        x = tf.strings.regex_replace(x,"ü", "u")
        x = tf.strings.regex_replace(x,"Á", "a")
        x = tf.strings.regex_replace(x,"É", "e")
        x = tf.strings.regex_replace(x,"Í", "i")
        x = tf.strings.regex_replace(x,"Ó", "o")
        x = tf.strings.regex_replace(x,"Ú", "u")
        x = tf.strings.regex_replace(x,"Ü", "u")
        x = tf.strings.regex_replace(x,"Ü", "u")
        x = tf.strings.regex_replace(x,"[?¿¡!@#$-_\?+¿{}*/]", "")
        x = tf.strings.regex_replace(x," +", " ")
        x = tf.strings.strip(x)
        x = tf.strings.split(x)
        x = self.table.lookup(x)
        x_as_vector = tf.reshape(x, [-1])
        zero_padding = tf.zeros([191] - tf.shape(x_as_vector), dtype=x.dtype)
        x = tf.concat([x_as_vector, zero_padding], 0)
        return x
        

    def call(self, inputs):
        x = tf.map_fn(lambda tx: self.sentences_transform(tx), elems = inputs,dtype=tf.int64)
        return x

    def get_config(self):
        return {'word_index': self.word_index}

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

with open(<tokenizer_path>) as f:
    data = json.load(f)
    tokenizer = tokenizer_from_json(data)

moderator = load_model(<final model path ('.h5')>)
word_index = tokenizer.word_index
text_bytes = tf.keras.Input(shape=(), name='image_bytes', dtype=tf.string)
x = VocabLookup(word_index)(text_bytes)
output = moderator(x)
model = tf.keras.models.Model(text_bytes, output)

如果你做总结，你会得到这样的东西：

model.summary()
Model: "functional_57"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
image_bytes (InputLayer)     [(None,)]                 0         
_________________________________________________________________
vocab_lookup_60 (VocabLookup (None, None)              0         
_________________________________________________________________
sequential_1 (Sequential)    (None, 1)                 1354369   
=================================================================
Total params: 1,354,369
Trainable params: 1,354,369
Non-trainable params: 0

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

save_path = <your_serving_model_path>
tf.saved_model.save(model,  save_path)

推荐阅读