tensorflow-gpu 2.0.0版本下keras.layers.StringLookup替代方案求问

keras.layers.StringLookup为TensorFlow 2.1之后才引入的内置层，TF2.0.0版本无对应实现，可采用以下两种兼容方案：

方案一：自定义等价功能Keras层（可嵌入模型计算图） #

该方案可实现与原生StringLookup基本一致的功能，支持配置词汇表、OOV处理规则，可直接嵌入模型计算图用于推理阶段的字符串输入处理：

import tensorflow as tf
from tensorflow.keras.layers import Layer

class CustomStringLookup(Layer):
    def __init__(self, vocabulary, oov_token="[OOV]", mask_token=None, **kwargs):
        super(CustomStringLookup, self).__init__(**kwargs)
        self.vocabulary = vocabulary
        self.oov_token = oov_token
        self.mask_token = mask_token
        # 构建静态哈希映射表
        self.lookup_table = tf.lookup.StaticHashTable(
            initializer=tf.lookup.KeyValueTensorInitializer(
                keys=tf.constant(vocabulary),
                values=tf.constant(list(range(len(vocabulary))), dtype=tf.int64)
            ),
            # OOV索引默认设置为词汇表长度，和原生接口逻辑对齐
            default_value=tf.constant(len(vocabulary), dtype=tf.int64)
        )

    def call(self, inputs):
        return self.lookup_table.lookup(inputs)

    def get_config(self):
        config = super(CustomStringLookup, self).get_config()
        config.update({
            "vocabulary": self.vocabulary,
            "oov_token": self.oov_token,
            "mask_token": self.mask_token
        })
        return config

使用示例：

# 自定义词汇表
vocab = ["苹果", "香蕉", "橙子", "葡萄"]
lookup_layer = CustomStringLookup(vocabulary=vocab)
# 测试输入
test_input = tf.constant(["苹果", "橙子", "西瓜", "香蕉"])
print(lookup_layer(test_input))
# 输出：tf.Tensor([0 2 4 1], shape=(4,), dtype=int64)
# 其中值4对应OOV词汇的索引

方案二：离线预处理映射（适合训练前可获取全量语料场景） #

如果不需要将字符串映射逻辑嵌入模型计算图，可在训练前用tf.keras.preprocessing.text.Tokenizer完成全量文本的离线编码，直接将编码后的整数序列输入模型：

from tensorflow.keras.preprocessing.text import Tokenizer

# 全量训练语料
corpus = ["苹果 香蕉 橙子", "橙子 葡萄 苹果", "香蕉 葡萄 西瓜"]
# 初始化分词器，oov_token指定未登录词标记，索引1为OOV，0保留为填充位
tokenizer = Tokenizer(oov_token="[OOV]")
tokenizer.fit_on_texts(corpus)
# 文本转整数序列
encoded_seq = tokenizer.texts_to_sequences(["苹果 橙子 西瓜 草莓"])
print(encoded_seq)
# 输出：[[2, 4, 6, 1]]

原始报错参考 #

内容的提问来源于stack exchange，提问作者Kunal Kamble