RSVP for your your local TensorFlow Everywhere event today!

tf.estimator.tpu.experimental.EmbeddingConfigSpec

View source on GitHub

Class to keep track of the specification for TPU embeddings.

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>