Type tf.tpu.experimental
Namespace tensorflow
Methods
Properties
Public static methods
_TPUEmbeddingColumnV2 embedding_column(CategoricalColumn categorical_column, int dimension, string combiner, object initializer, int max_sequence_length, PythonFunctionContainer learning_rate_fn)
TPU version of `tf.compat.v1.feature_column.embedding_column`. Note that the interface for
tf.tpu.experimental.embedding_column
is
different from that of `tf.compat.v1.feature_column.embedding_column`: The
following arguments are NOT supported: `ckpt_to_load_from`,
`tensor_name_in_ckpt`, `max_norm` and `trainable`. Use this function in place of `tf.compat.v1.feature_column.embedding_column`
when you want to use the TPU to accelerate your embedding lookups via TPU
embeddings. ```
column = tf.feature_column.categorical_column_with_identity(...)
tpu_column = tf.tpu.experimental.embedding_column(column, 10)
...
def model_fn(features):
dense_feature = tf.keras.layers.DenseFeature(tpu_column)
embedded_feature = dense_feature(features)
... estimator = tf.estimator.tpu.TPUEstimator(
model_fn=model_fn,
...
embedding_config_spec=tf.estimator.tpu.experimental.EmbeddingConfigSpec(
column=[tpu_column],
...))
```
Parameters
-
CategoricalColumn
categorical_column - A categorical column returned from `categorical_column_with_identity`, `weighted_categorical_column`, `categorical_column_with_vocabulary_file`, `categorical_column_with_vocabulary_list`, `sequence_categorical_column_with_identity`, `sequence_categorical_column_with_vocabulary_file`, `sequence_categorical_column_with_vocabulary_list`
-
int
dimension - An integer specifying dimension of the embedding, must be > 0.
-
string
combiner - A string specifying how to reduce if there are multiple entries
in a single row for a non-sequence column. For more information, see
tf.feature_column.embedding_column
. -
object
initializer - A variable initializer function to be used in embedding variable initialization. If not specified, defaults to `tf.compat.v1.truncated_normal_initializer` with mean `0.0` and standard deviation `1/sqrt(dimension)`.
-
int
max_sequence_length - An non-negative integer specifying the max sequence length. Any sequence shorter then this will be padded with 0 embeddings and any sequence longer will be truncated. This must be positive for sequence features and 0 for non-sequence features.
-
PythonFunctionContainer
learning_rate_fn - A function that takes global step and returns learning rate for the embedding table.
Returns
-
_TPUEmbeddingColumnV2
- A `_TPUEmbeddingColumnV2`.
Topology initialize_tpu_system(TPUClusterResolver cluster_resolver)
Initialize the TPU devices.
Parameters
-
TPUClusterResolver
cluster_resolver - A tf.distribute.cluster_resolver.TPUClusterResolver, which provides information about the TPU cluster.
Returns
-
Topology
- The tf.tpu.Topology object for the topology of the TPU cluster.
object initialize_tpu_system_dyn(object cluster_resolver)
Initialize the TPU devices.
Parameters
-
object
cluster_resolver - A tf.distribute.cluster_resolver.TPUClusterResolver, which provides information about the TPU cluster.
Returns
-
object
- The tf.tpu.Topology object for the topology of the TPU cluster.