如何将 tf.lookup 表与 TensorFlow 2.0 keras 和 MLFlow 结合使用

Question

我花了大约 5 个小时左右的时间尝试让 TF 2.0 keras API 与 tf.lookup API 一起工作。我的训练脚本还使用 DataBricks 和mlflow.keras. MLFlow 要求模型被序列化，我认为这就是给我带来问题的原因。问题是：如何将 tf.lookup 表与 TensorFlow 2.0 keras 模型 API 和 MLFlow 结合使用。

当我尝试直接将功能性 Keras API 与 table.lookup 一起使用时，我遇到了 keras 序列化问题：

table = tf.lookup.StaticVocabularyTable(tf.lookup.TextFileInitializer(vocab_path, tf.string, 0, tf.int64, 1, delimiter=","), 1)
categorical_indices = table.lookup(categorical_input)

将上述调用包装在一层中tf.keras.layers.Lambda没有帮助。我收到与资源句柄或缺少tf变量相关的错误...

Answer 1

在这里分享解决方案以减轻其他人的痛苦。这是我发现有效的解决方案：

vocab_path = os.path.join(mount_point, 'category_vocab.csv')

class VocabLookup(layers.Layer):
  def __init__(self, vocab_path, num_oov_buckets, **kwargs):
    self.vocab_path = vocab_path
    self.num_oov_buckets = num_oov_buckets
    super(VocabLookup, self).__init__(**kwargs)

  def build(self, input_shape):

    vocab_initializer = tf.lookup.TextFileInitializer(
      self.vocab_path, tf.string, 0, tf.int64, 1, delimiter=",")
    self.table = tf.lookup.StaticVocabularyTable(vocab_initializer, self.num_oov_buckets)

    self.built = True

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

  def get_config(self):
    return {'vocab_path': self.vocab_path, 'num_oov_buckets': self.num_oov_buckets}

lookup_table = VocabLookup(vocab_path, 1)

categorical_indices = lookup_table(categorical_input)

基本上，如果您引用任何外部变量（包括 tf 或tensorflow 模块），请不要使用layers.Lambda。例如，这对我不起作用：

def reduce_sum(x):
  return tf.reduce_sum(x, axis=1)

embedding_sum = layers.Lambda(reduce_sum)

categorical_features = embedding_sum(categorical_embeddings)

但这有效：

class ReduceSum(layers.Layer):
  def call(self, inputs):
    return tf.reduce_sum(inputs, axis=1)

embedding_sum = ReduceSum()
categorical_features = embedding_sum(categorical_embeddings)

他们layers.Lambda似乎不喜欢升值。