tf.tpu.experimental.embedding.serving_embedding_lookup  |  TensorFlow v2.16.1 (original) (raw)

tf.tpu.experimental.embedding.serving_embedding_lookup

Stay organized with collections Save and categorize content based on your preferences.

Apply standard lookup ops with tf.tpu.experimental.embedding configs.

View aliases

Compat aliases for migration

SeeMigration guide for more details.

tf.compat.v1.tpu.experimental.embedding.serving_embedding_lookup

tf.tpu.experimental.embedding.serving_embedding_lookup(
    inputs: Any,
    weights: Optional[Any],
    tables: Dict[tf.tpu.experimental.embedding.TableConfig, tf.Variable],
    feature_config: Union[tf.tpu.experimental.embedding.FeatureConfig, Iterable]
) -> Any

This function is a utility which allows using thetf.tpu.experimental.embedding config objects with standard lookup functions. This can be used when exporting a model which usestf.tpu.experimental.embedding.TPUEmbedding for serving on CPU. In particulartf.tpu.experimental.embedding.TPUEmbedding only supports lookups on TPUs and should not be part of your serving graph.

Note that TPU specific options (such as max_sequence_length) in the configuration objects will be ignored.

In the following example we take a trained model (see the documentation fortf.tpu.experimental.embedding.TPUEmbedding for the context) and create a saved model with a serving function that will perform the embedding lookup and pass the results to your model:

model = model_fn(...)
embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=1024,
    optimizer=tf.tpu.experimental.embedding.SGD(0.1))
checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
checkpoint.restore(...)

@tf.function(input_signature=[{'feature_one': tf.TensorSpec(...),
                               'feature_two': tf.TensorSpec(...),
                               'feature_three': tf.TensorSpec(...)}])
def serve_tensors(embedding_features):
  embedded_features = tf.tpu.experimental.embedding.serving_embedding_lookup(
      embedding_features, None, embedding.embedding_tables,
      feature_config)
  return model(embedded_features)

model.embedding_api = embedding
tf.saved_model.save(model,
                    export_dir=...,
                    signatures={'serving_default': serve_tensors})