tf.estimator.tpu.experimental.EmbeddingConfigSpec

View source on GitHub

Class to keep track of the specification for TPU embeddings.

tf.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))))

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>

<tr>
<td>
`feature_columns`
</td>
<td>
All embedding `FeatureColumn`s used by model.
</td>
</tr><tr>
<td>
`optimization_parameters`
</td>
<td>
An instance of `AdagradParameters`,
`AdamParameters` or `StochasticGradientDescentParameters`. This
optimizer will be applied to all embedding variables specified by
`feature_columns`.
</td>
</tr><tr>
<td>
`clipping_limit`
</td>
<td>
(Optional) Clipping limit (absolute value).
</td>
</tr><tr>
<td>
`pipeline_execution_with_tensor_core`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`experimental_gradient_multiplier_fn`
</td>
<td>
(Optional) A Fn taking global step as
input returning the current multiplier for all embedding gradients.
</td>
</tr><tr>
<td>
`feature_to_config_dict`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`table_to_config_dict`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`partition_strategy`
</td>
<td>
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.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>

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





<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>

<tr>
<td>
`feature_columns`
</td>
<td>

</td>
</tr><tr>
<td>
`optimization_parameters`
</td>
<td>

</td>
</tr><tr>
<td>
`clipping_limit`
</td>
<td>

</td>
</tr><tr>
<td>
`pipeline_execution_with_tensor_core`
</td>
<td>

</td>
</tr><tr>
<td>
`experimental_gradient_multiplier_fn`
</td>
<td>

</td>
</tr><tr>
<td>
`feature_to_config_dict`
</td>
<td>

</td>
</tr><tr>
<td>
`table_to_config_dict`
</td>
<td>

</td>
</tr><tr>
<td>
`partition_strategy`
</td>
<td>

</td>
</tr>
</table>