tf.tpu.experimental.embedding.serving_embedding_lookup

Применить стандартные операции поиска с конфигурациями tf.tpu.experimental.embedding.

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

Эта функция — утилита, позволяющая использовать объекты конфигурации tf.tpu.experimental.embedding со стандартными функциями поиска. Это может быть полезно при экспорте модели, использующей tf.tpu.experimental.embedding.TPUEmbedding для обслуживания на процессоре. В частности, tf.tpu.experimental.embedding.TPUEmbedding поддерживает поиск только на TPUs и не должна входить в вашу служебную графику.

Обратите внимание, что специфичные для TPU параметры (например, max_sequence_length) в объектах конфигурации будут игнорироваться.

В следующем примере мы возьмем обученную модель (см. документацию для tf.tpu.experimental.embedding.TPUEmbedding для контекста) и создадим сохранённую модель с функцией обслуживания, которая выполнит поиск по вложениям и передаст результаты вашей модели:

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})

Примечание: Важно назначить объект API вложений члену вашей модели, так как tf.saved_model.save поддерживает сохранение переменных только как один объект Trackable. Поскольку веса модели находятся в model, а таблица вложений управляется embedding, мы назначаем embedding атрибуту model, чтобы tf.saved_model.save мог найти переменные вложений.

Примечание: Та же самая serve_tensors функция и вызов tf.saved_model.save будут работать непосредственно из обучения.