View source on GitHub |
The TPUEmbedding mid level API running on TPU with sparse core accelerator.
tf.tpu.experimental.embedding.TPUEmbeddingV2(
feature_config: Union[tf.tpu.experimental.embedding.FeatureConfig
, Iterable],
optimizer: Optional[tpu_embedding_v2_utils._Optimizer] = None,
pipeline_execution_with_tensor_core: bool = False
)
Args | |
---|---|
feature_config
|
A nested structure of
tf.tpu.experimental.embedding.FeatureConfig configs.
|
optimizer
|
An instance of one of tf.tpu.experimental.embedding.SGD ,
tf.tpu.experimental.embedding.Adagrad or
tf.tpu.experimental.embedding.Adam . When not created under TPUStrategy
may be set to None to avoid the creation of the optimizer slot
variables, useful for optimizing memory consumption when exporting the
model for serving where slot variables aren't needed.
|
pipeline_execution_with_tensor_core
|
If True, the TPU embedding computations will overlap with the TensorCore computations (and hence will be one step old). Set to True for improved performance. |
Raises | |
---|---|
ValueError
|
If optimizer is not one of tf.tpu.experimental.embedding.(SGD, Adam or Adagrad) or None when created under a TPUStrategy. |
RuntimeError
|
If not created under TPUStrategy. |
Methods
apply_gradients
apply_gradients(
gradients: Any, preserved_outputs: Dict[str, PartitionedCsrFormatTensor]
)
Applies the gradient update to the embedding tables.
If a gradient of None
is passed in any position of the nested structure,
then a gradient update with a zero gradient is applied for that feature.
For optimizers like SGD or Adagrad, this is the same as applying no update
at all. For lazy Adam and other sparsely applied optimizers with decay,
ensure you understand the effect of applying a zero gradient.
Args | |
---|---|
gradients
|
A nested structure of gradients, with structure matching the
feature_config passed to this object.
|
preserved_outputs
|
A dicts of PartitionedCsrFormatTensor, coming from the second output of the embedding lookup call. |
Raises | |
---|---|
RuntimeError
|
if not built. |
ValueError
|
If a non-tf.Tensor non-None gradient is passed in, or a
tf.Tensor of the incorrect shape is passed in. Also if
the size of any sequence in gradients does not match corresponding
sequence in feature_config .
|
TypeError
|
If the type of any sequence in gradients does not match
corresponding sequence in feature_config .
|
build
build()
Create variables and slots variables for TPU embeddings.
dequeue
dequeue(
partitioned_tensors: Tuple[Dict[str, PartitionedCsrFormatTensor], int, int]
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]
Perform embedding lookup.
embedding_lookup
embedding_lookup(
features: Any, weights: Optional[Any] = None
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]
Perform embedding lookup on the input feature.
Args | |
---|---|
features
|
A nested structure of tf.Tensor s, tf.SparseTensor s or
tf.RaggedTensor s, with the same structure as feature_config . Inputs
will be downcast to tf.int32 . Only one type out of tf.SparseTensor
or tf.RaggedTensor is supported per call.
|
weights
|
If not None , a nested structure of tf.Tensor s,
tf.SparseTensor s or tf.RaggedTensor s, matching the above, except
that the tensors should be of float type (and they will be downcast to
tf.float32 ). For tf.SparseTensor s we assume the indices are the
same for the parallel entries from features and similarly for
tf.RaggedTensor s we assume the row_splits are the same.
|
Raises | |
---|---|
ValueError
|
If the input feature is not one of the Tensor, SparseTensor or RaggedTensor type. |
TypeError
|
If the type of any sequence in features does not match
corresponding sequence in feature_config . Similarly for weights , if
not None .
|
Returns | |
---|---|
packed_activations
|
Embedding lookup results packed as the same sequence of the input feature. |
packed_output
|
A dict of PartitionedCsrFormatTensors. |
enqueue
enqueue(
features: Any, weights: Optional[Any] = None, device: Optional[str] = None
) -> Any
Preprocessing the features on host.
preprocess_features
@staticmethod
preprocess_features( num_replicas_in_sync: int, max_ids_per_chip_per_sample: int, max_minibatches_per_sc: int, num_sc_per_chip: int, num_sc_shards: int, stacked_table_to_tables: Dict[str, Any], table_to_stacked_table_offset: Dict[str, Tuple[str, int, int]], table_to_sample_count: Dict[str, int], feature_to_sample_offset: Dict[str, int], flat_features: Any, flat_inputs: Any, flat_weights: Optional[Any] = None ) -> Any
Function to preprocess features.
__call__
__call__(
features: Any, weights: Optional[Any] = None
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]
Call the mid level api to do embedding lookup.