tf.tpu.experimental.embedding_column

TPU версия tf.compat.v1.feature_column.embedding_column.

tf.tpu.experimental.embedding_column(
    categorical_column, dimension, combiner='mean', initializer=None,
    max_sequence_length=0, learning_rate_fn=None
)

Обратите внимание, что интерфейс tf.tpu.experimental.embedding_column отличается от интерфейса tf.compat.v1.feature_column.embedding_column: следующие аргументы НЕ поддерживаются: ckpt_to_load_from, tensor_name_in_ckpt, max_norm и trainable.

Используйте эту функцию вместо tf.compat.v1.feature_column.embedding_column, когда требуется использовать TPU для ускорения поиска встраиваемых значений с помощью TPU встраивания.

column = tf.feature_column.categorical_column_with_identity(...)
tpu_column = tf.tpu.experimental.embedding_column(column, 10)
...
def model_fn(features):
  dense_feature = tf.keras.layers.DenseFeature(tpu_column)
  embedded_feature = dense_feature(features)
  ...

estimator = tf.estimator.tpu.TPUEstimator(
    model_fn=model_fn,
    ...
    embedding_config_spec=tf.estimator.tpu.experimental.EmbeddingConfigSpec(
      column=[tpu_column],
      ...))