tf.tpu.experimental.embedding.serving_embedding_lookup

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

tf.tpu.experimental.embedding.serving_embedding_lookup(
    inputs, weights, tables, feature_config
)

Эта функция является утилитой, которая позволяет использовать объекты конфигурации 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_featurese):
  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 будут работать напрямую из обучения.