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, для обслуживания на CPU. В частности, tf.tpu.experimental.embedding.TPUEmbedding поддерживает поиск только на TPU и не должна входить в граф обслуживания.

Обратите внимание, что параметры, относящиеся к 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 будут работать непосредственно во время обучения.