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)))) <!-- 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> `tensor_core_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>
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.3/api_docs/python/tf/compat/v1/estimator/tpu/experimental/EmbeddingConfigSpec