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

View source on GitHub

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

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:


* <b>`feature_columns`</b>: All embedding `FeatureColumn`s used by model.
* <b>`optimization_parameters`</b>: An instance of `AdagradParameters`,
  `AdamParameters` or `StochasticGradientDescentParameters`. This
  optimizer will be applied to all embedding variables specified by
  `feature_columns`.
* <b>`clipping_limit`</b>: (Optional) Clipping limit (absolute value).
* <b>`pipeline_execution_with_tensor_core`</b>: 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.
* <b>`experimental_gradient_multiplier_fn`</b>: (Optional) A Fn taking global step as
  input returning the current multiplier for all embedding gradients.
* <b>`feature_to_config_dict`</b>: A dictionary mapping features 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.
* <b>`table_to_config_dict`</b>: A dictionary mapping features 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.
* <b>`partition_strategy`</b>: A string, determining how tensors are sharded to the
  tpu hosts. See <a href="../../../../../../tf/nn/safe_embedding_lookup_sparse"><code>tf.nn.safe_embedding_lookup_sparse</code></a> 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.


#### Attributes:

* <b>`feature_columns`</b>
* <b>`optimization_parameters`</b>
* <b>`clipping_limit`</b>
* <b>`pipeline_execution_with_tensor_core`</b>
* <b>`experimental_gradient_multiplier_fn`</b>
* <b>`feature_to_config_dict`</b>
* <b>`table_to_config_dict`</b>
* <b>`partition_strategy`</b>


#### Raises:


* <b>`ValueError`</b>: If the feature_columns are not specified.
* <b>`TypeError`</b>: If the feature columns are not of ths correct type (one of
  _SUPPORTED_FEATURE_COLUMNS, _TPU_EMBEDDING_COLUMN_CLASSES OR
  _EMBEDDING_COLUMN_CLASSES).
* <b>`ValueError`</b>: If `optimization_parameters` is not one of the required types.