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.md"><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.