tf.compat.v1.estimator.tpu.experimental.EmbeddingConfigSpec

Class to keep track of the specification for TPU embeddings.

tf.compat.v1.estimator.tpu.experimental.EmbeddingConfigSpec(
    feature_columns=None, optimization_parameters=None, clipping_limit=None,
    pipeline_execution_with_tensor_core=False,
    experimental_gradient_multiplier_fn=None, feature_to_config_dict=None,
    table_to_config_dict=None, partition_strategy='div',
    profile_data_directory=None
)

Pass this class to tf.estimator.tpu.TPUEstimator via the embedding_config_spec parameter. At minimum you need to specify feature_columns and optimization_parameters. The feature columns passed should be created with some combination of tf.tpu.experimental.embedding_column and tf.tpu.experimental.shared_embedding_columns.

TPU embeddings do not support arbitrary Tensorflow optimizers and the main optimizer you use for your model will be ignored for the embedding table variables. Instead TPU embeddigns support a fixed set of predefined optimizers that you can select from and set the parameters of. These include adagrad, adam and stochastic gradient descent. Each supported optimizer has a Parameters class in the tf.tpu.experimental namespace.

column_a = tf.feature_column.categorical_column_with_identity(...)
column_b = tf.feature_column.categorical_column_with_identity(...)
column_c = tf.feature_column.categorical_column_with_identity(...)
tpu_shared_columns = tf.tpu.experimental.shared_embedding_columns(
    [column_a, column_b], 10)
tpu_non_shared_column = tf.tpu.experimental.embedding_column(
    column_c, 10)
tpu_columns = [tpu_non_shared_column] + tpu_shared_columns
...
def model_fn(features):
  dense_features = tf.keras.layers.DenseFeature(tpu_columns)
  embedded_feature = dense_features(features)
  ...

estimator = tf.estimator.tpu.TPUEstimator(
    model_fn=model_fn,
    ...
    embedding_config_spec=tf.estimator.tpu.experimental.EmbeddingConfigSpec(
        column=tpu_columns,
        optimization_parameters=(
            tf.estimator.tpu.experimental.AdagradParameters(0.1))))

Args
feature_columns	All embedding FeatureColumns used by model.
optimization_parameters	An instance of AdagradParameters, AdamParameters or StochasticGradientDescentParameters. This optimizer will be applied to all embedding variables specified by feature_columns.
clipping_limit	(Optional) Clipping limit (absolute value).
pipeline_execution_with_tensor_core	setting this to True makes training faster, but trained model will be different if step N and step N+1 involve the same set of embedding IDs. Please see tpu_embedding_configuration.proto for details.
experimental_gradient_multiplier_fn	(Optional) A Fn taking global step as input returning the current multiplier for all embedding gradients.
feature_to_config_dict	A dictionary mapping feature names to instances of the class FeatureConfig. Either features_columns or the pair of feature_to_config_dict and table_to_config_dict must be specified.
table_to_config_dict	A dictionary mapping feature names to instances of the class TableConfig. Either features_columns or the pair of feature_to_config_dict and table_to_config_dict must be specified.
partition_strategy	A string, determining how tensors are sharded to the tpu hosts. See tf.nn.safe_embedding_lookup_sparse for more details. Allowed value are "div" and "mod"'. If"mod"is used, evaluation and exporting the model to CPU will not work as expected. </td> </tr><tr> <td>profile_data_directory`	Directory where embedding lookup statistics are stored. These statistics summarize information about the inputs to the embedding lookup operation, in particular, the average number of embedding IDs per example and how well the embedding IDs are load balanced across the system. The lookup statistics are used during TPU initialization for embedding table partitioning. Collection of lookup statistics is done at runtime by profiling the embedding inputs: only 3% of input samples are profiled to minimize host CPU overhead. Once a suitable number of samples are profiled, the lookup statistics are saved to table-specific files in the profile data directory generally at the end of a TPU training loop. The filename corresponding to each table is obtained by hashing table specific parameters (e.g., table name and number of features) and global configuration parameters (e.g., sharding strategy and task count). The same profile data directory can be shared among several models to reuse embedding lookup statistics.

Raises
ValueError	If the feature_columns are not specified.
TypeError	If the feature columns are not of ths correct type (one of _SUPPORTED_FEATURE_COLUMNS, _TPU_EMBEDDING_COLUMN_CLASSES OR _EMBEDDING_COLUMN_CLASSES).
ValueError	If optimization_parameters is not one of the required types.

Attributes
feature_columns
tensor_core_feature_columns
optimization_parameters
clipping_limit
pipeline_execution_with_tensor_core
experimental_gradient_multiplier_fn
feature_to_config_dict
table_to_config_dict
partition_strategy
profile_data_directory