wehub-resource-sync
446 lines
18 KiB
Python
446 lines
18 KiB
Python
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Mid level API for TPU Embeddings without Embedding Accelerator."""
|
|
|
|
from typing import Any, Dict, Iterable, Optional, Text, Union
|
|
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.distribute import tpu_strategy
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import sparse_tensor
|
|
from tensorflow.python.framework import tensor
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import embedding_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import sparse_ops
|
|
from tensorflow.python.ops import variables as tf_variables
|
|
from tensorflow.python.ops.ragged import ragged_tensor
|
|
from tensorflow.python.tpu import tpu_embedding_base
|
|
from tensorflow.python.tpu import tpu_embedding_v2_utils
|
|
from tensorflow.python.tpu import tpu_replication
|
|
from tensorflow.python.util import nest
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
|
|
@tf_export("tpu.experimental.embedding.TPUEmbeddingV0")
|
|
class TPUEmbeddingV0(tpu_embedding_base.TPUEmbeddingBase):
|
|
"""The TPUEmbedding mid level API running on TPU without Embedding accelerator.
|
|
|
|
NOTE: This mid level API is not intended for large embedding table lookup.
|
|
Embedding tables will be replicated across devices rather than sharding
|
|
across them. To do large embedding table lookup, please use the
|
|
`tpu.experimental.embedding.TPUEmbedding` class. This class is an alternative
|
|
way to do embedding lookups when the TPU doesn't support any version of
|
|
embedding feature. See
|
|
`tpu.experimental.tpu_hardware_feature.embedding_feature` for a detailed
|
|
explanation.
|
|
|
|
This class has to be created under the `TPUStrategy`, Otherwise a RuntimeError
|
|
will be raised.
|
|
```python
|
|
strategy = tf.distribute.TPUStrategy(...)
|
|
with strategy.scope():
|
|
embedding = tf.tpu.experimental.embedding.TPUEmbeddingV0(
|
|
feature_config=feature_config,
|
|
optimizer=tf.tpu.experimental.embedding.SGD(0.1))
|
|
```
|
|
When creating a distributed dataset that is to be passed to the lookup
|
|
operation a special input option must be specified:
|
|
|
|
```python
|
|
distributed_dataset = (
|
|
strategy.distribute_datasets_from_function(
|
|
dataset_fn=...,
|
|
options=tf.distribute.InputOptions(
|
|
experimental_fetch_to_device=False))
|
|
dataset_iterator = iter(distributed_dataset)
|
|
```
|
|
|
|
Below is an example of a training and evaluation step:
|
|
|
|
```python
|
|
optimizer = tf.keras.optimizers.SGD(0.1)
|
|
|
|
@tf.function
|
|
def training_step(dataset_iterator, num_steps):
|
|
def tpu_step(embedding_features):
|
|
with tf.GradientTape() as tape:
|
|
tape.watch(embedding.embedding_table.values())
|
|
activation = embedding(embedding_features)
|
|
model_output = model(activations)
|
|
loss = ... # some function of labels and model_output
|
|
|
|
embedding_gradients = tape.gradient(loss,
|
|
embedding.embedding_table.values())
|
|
optimizer.apply_gradients(list(zip(gradients,
|
|
mid_level_api.embedding_tables.values())))
|
|
# Insert your model gradient and optimizer application here
|
|
|
|
for _ in tf.range(num_steps):
|
|
strategy.run(tpu_step, args=(next(dataset_iterator), ))
|
|
|
|
@tf.function
|
|
def evalution_step(dataset_iterator, num_steps):
|
|
def tpu_step(embedding_features):
|
|
activations = embedding(embedding_features)
|
|
model_output = model(activations)
|
|
# Insert your evaluation code here.
|
|
|
|
for _ in tf.range(num_steps):
|
|
strategy.run(tpu_step, args=(next(dataset_iterator), ))
|
|
```
|
|
|
|
NOTE: The optimizer used here is a Keras optimizer. In order to make the slot
|
|
variable creation stay consistent between Keras optimizers and
|
|
embedding optimizers, the `slot_variable_creation_fn` argument of the
|
|
embedding optimizers has to be passed with the Keras `add_slot` function. Also
|
|
note that the slot names might be slightly different between them.
|
|
|
|
```python
|
|
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)
|
|
|
|
def slot_variable_creation_fn(table, slot_names, slot_initializers):
|
|
slots = {}
|
|
for slot, initializer in zip(slot_names, slot_initializers):
|
|
slots[slot] = optimizer.add_slot(table, slot, initializer)
|
|
return slots
|
|
|
|
embedding_optimizer = tf.experimental.embedding.Adagrad(
|
|
learning_rate=0.1,
|
|
slot_variable_creation_fn=slot_variable_creation_fn)
|
|
|
|
# Use the embedding optimizer to create mid level api and keras optimizer to
|
|
# apply gradients.
|
|
```
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable], # pylint:disable=g-bare-generic
|
|
optimizer: Optional[tpu_embedding_v2_utils._Optimizer]): # pylint:disable=protected-access
|
|
super(TPUEmbeddingV0, self).__init__(feature_config, optimizer)
|
|
self._strategy = distribute_lib.get_strategy()
|
|
if not isinstance(self._strategy,
|
|
(tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2)):
|
|
raise RuntimeError(
|
|
"TPUEmbeddingV0 should be created under TPUStrategy but found {}."
|
|
.format(self._strategy))
|
|
self._built = False
|
|
|
|
@property
|
|
def embedding_tables(
|
|
self) -> Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable]:
|
|
"""Returns a dict of embedding tables, keyed by `TableConfig`."""
|
|
self._maybe_build()
|
|
# Only return the tables and not the slot variables.
|
|
return {
|
|
table: self._variables[table.name]["parameters"]
|
|
for table in self._table_config
|
|
}
|
|
|
|
def _create_variables_and_slots(
|
|
self) -> Dict[Text, Dict[Text, tf_variables.Variable]]:
|
|
"""Create variables for TPU embeddings.
|
|
|
|
Note that this will always ensure that the variable is created under the
|
|
TPUStrategy.
|
|
|
|
Returns:
|
|
A dict of dicts. The outer dict is keyed by the table names and the inner
|
|
dicts are keyed by 'parameters' and the slot variable names.
|
|
"""
|
|
variables = {}
|
|
for table in self._table_config:
|
|
# created TPUDistributedVariable.
|
|
variables[table.name] = self._create_variables(table, trainable=True)
|
|
return variables
|
|
|
|
def _maybe_build(self):
|
|
if not self._built:
|
|
# This can be called while tracing a function, so we wrap the
|
|
# initialization code with init_scope so it runs eagerly, this means that
|
|
# it will not be included in the function graph generated by tracing so
|
|
# that we can be sure that we only initialize the TPU for embeddings
|
|
# exactly once.
|
|
with ops.init_scope():
|
|
self.build()
|
|
|
|
def _apply_combiner_to_embeddings(
|
|
self,
|
|
embeddings: tensor.Tensor,
|
|
weight: tensor.Tensor,
|
|
combiner: Optional[Text] = None) -> tensor.Tensor:
|
|
"""Apply the combiner to the embedding look up result on second to last axis.
|
|
|
|
Args:
|
|
embeddings: A Tensor of the embedding lookup result.
|
|
weight: A Tensor of weight which has the same shape of the embeddings.
|
|
combiner: One of "mean", "sum", "sqrtn". Defaults to "mean".
|
|
|
|
Raises:
|
|
ValueError: If the combiner is not one of 'mean', 'sqrtn' or 'sum'.
|
|
Returns:
|
|
A Tensor.
|
|
"""
|
|
if combiner is None:
|
|
combiner = "mean"
|
|
if combiner == "sum":
|
|
embeddings = math_ops.reduce_sum(embeddings, axis=-2)
|
|
elif combiner == "mean":
|
|
embeddings = math_ops.reduce_sum(embeddings, axis=-2)
|
|
weight_sum = math_ops.reduce_sum(weight, axis=-2)
|
|
embeddings = math_ops.div_no_nan(embeddings, weight_sum)
|
|
elif combiner == "sqrtn":
|
|
embeddings = math_ops.reduce_sum(embeddings, axis=-2)
|
|
weight_squared = math_ops.pow(weight, 2)
|
|
weight_sum = math_ops.reduce_sum(weight_squared, axis=-2)
|
|
weight_sum_sqrt = math_ops.sqrt(weight_sum)
|
|
embeddings = math_ops.div_no_nan(embeddings, weight_sum_sqrt)
|
|
else:
|
|
raise ValueError(
|
|
f"combiner must be one of 'mean', 'sqrtn' or 'sum', got {combiner}")
|
|
return embeddings
|
|
|
|
def _pad_or_truncate_with_sequence_length(
|
|
self, embeddings: tensor.Tensor, sequence_length: int
|
|
) -> tensor.Tensor:
|
|
"""Pad or truncate the embedding lookup result based on the sequence length.
|
|
|
|
Args:
|
|
embeddings: A rank 3 Tensor of the embedding lookup result.
|
|
sequence_length: number of the max sequence length set in the feature
|
|
config.
|
|
|
|
Returns:
|
|
A Tensor with second last axis padded or truncated.
|
|
"""
|
|
original_sequence_length = embeddings.shape[1]
|
|
if original_sequence_length > sequence_length:
|
|
embeddings = array_ops.slice(
|
|
embeddings, begin=[0, 0, 0], size=[-1, sequence_length, -1])
|
|
else:
|
|
embeddings = array_ops.pad(
|
|
embeddings,
|
|
paddings=[[0, 0], [0, sequence_length - original_sequence_length],
|
|
[0, 0]])
|
|
return embeddings
|
|
|
|
def embedding_lookup(self,
|
|
features: Any,
|
|
weights: Optional[Any] = None) -> Any:
|
|
"""Apply embedding lookup on TPUs using Tensorcore.
|
|
|
|
Note that all the sparse and ragged tensors will be converted to dense
|
|
tensors on CPU and then passed to the TPU to do embedding look up. Large
|
|
embedding lookup is not supported by this API, use the TPUEmbedding mid
|
|
level api instead.
|
|
|
|
Args:
|
|
features: a nested structure of Tensors, SparseTensors or RaggedTensors.
|
|
weights: a nested structure of Tensors, SparseTensors or RaggedTensors or
|
|
None for no weights. If not None, structure must match that of inputs,
|
|
but entries are allowed to be None.
|
|
|
|
Returns:
|
|
A nested structure of Tensors with the same structure as inputs.
|
|
"""
|
|
if not self._built:
|
|
self.build()
|
|
nest.assert_same_structure(features, self._feature_config)
|
|
|
|
flat_inputs = nest.flatten(features)
|
|
flat_weights = [None] * len(flat_inputs)
|
|
if weights is not None:
|
|
nest.assert_same_structure(features, weights)
|
|
flat_weights = nest.flatten(weights)
|
|
flat_features = nest.flatten_with_joined_string_paths(self._feature_config)
|
|
|
|
outputs = []
|
|
for inp, weight, (path, feature) in zip(flat_inputs, flat_weights,
|
|
flat_features):
|
|
table = self.embedding_tables[feature.table]
|
|
|
|
if weight is not None:
|
|
if isinstance(inp, tensor.Tensor):
|
|
raise ValueError(
|
|
"Weight specified for {}, but input is dense.".format(path))
|
|
elif type(weight) is not type(inp):
|
|
raise ValueError(
|
|
"Weight for {} is of type {} but it does not match type of the "
|
|
"input which is {}.".format(path, type(weight), type(inp)))
|
|
elif feature.max_sequence_length > 0:
|
|
raise ValueError("Weight specified for {}, but this is a sequence "
|
|
"feature.".format(path))
|
|
|
|
if isinstance(inp, tensor.Tensor):
|
|
if feature.max_sequence_length > 0:
|
|
raise ValueError(
|
|
"Feature {} is a sequence feature but a dense tensor "
|
|
"was passed.".format(path))
|
|
outputs.append(embedding_ops.embedding_lookup_v2(table, inp))
|
|
|
|
elif isinstance(inp, sparse_tensor.SparseTensor):
|
|
outputs.append(
|
|
self._embedding_lookup_for_sparse_tensor(inp, weight, table,
|
|
feature))
|
|
elif isinstance(inp, ragged_tensor.RaggedTensor):
|
|
outputs.append(
|
|
self._embedding_lookup_for_ragged_tensor(inp, weight, table,
|
|
feature))
|
|
else:
|
|
raise ValueError("Input {} is type {}. Tensor, SparseTensor or "
|
|
"RaggedTensor expected.".format(path, type(inp)))
|
|
return nest.pack_sequence_as(self._feature_config, outputs)
|
|
|
|
def _embedding_lookup_for_sparse_tensor(
|
|
self, inp: sparse_tensor.SparseTensor,
|
|
weight: Optional[sparse_tensor.SparseTensor],
|
|
table: tf_variables.Variable,
|
|
feature: tpu_embedding_v2_utils.FeatureConfig) -> tensor.Tensor:
|
|
"""Embedding lookup for sparse tensor based on its feature config.
|
|
|
|
Args:
|
|
inp: a single SparseTensor input.
|
|
weight: None or SparseTensor which has the same shape of the input.
|
|
table: a table variable.
|
|
feature: a feature config.
|
|
|
|
Returns:
|
|
Embedding lookup result.
|
|
"""
|
|
|
|
# This computation needs to placed outside of tpu as the size of the
|
|
# indices and values can change for different batch which can cause
|
|
# the program to re-compile.
|
|
def sparse_to_dense_computation(inp, weight):
|
|
if weight is None:
|
|
weight = sparse_tensor.SparseTensor(
|
|
inp.indices,
|
|
array_ops.ones_like(inp.values, dtype=dtypes.float32),
|
|
dense_shape=inp.dense_shape)
|
|
# Pad the sparse tensor to be dense tensor.
|
|
inp = sparse_ops.sparse_tensor_to_dense(inp)
|
|
weight = sparse_ops.sparse_tensor_to_dense(weight)
|
|
return inp, weight
|
|
|
|
inp, weight = tpu_replication.outside_compilation(
|
|
sparse_to_dense_computation, inp=inp, weight=weight)
|
|
|
|
embeddings = embedding_ops.embedding_lookup_v2(table, inp)
|
|
weight = array_ops.expand_dims(weight, -1)
|
|
embeddings *= weight
|
|
if not feature.output_shape and feature.max_sequence_length > 0:
|
|
embeddings = self._pad_or_truncate_with_sequence_length(
|
|
embeddings, feature.max_sequence_length)
|
|
else:
|
|
embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
|
|
feature.table.combiner)
|
|
return embeddings
|
|
|
|
def _embedding_lookup_for_ragged_tensor(
|
|
self, inp: ragged_tensor.RaggedTensor,
|
|
weight: Optional[ragged_tensor.RaggedTensor],
|
|
table: tf_variables.Variable,
|
|
feature: tpu_embedding_v2_utils.FeatureConfig) -> tensor.Tensor:
|
|
"""Embedding lookup for ragged tensor based on its feature config.
|
|
|
|
Args:
|
|
inp: a single rank 2 RaggedTensor input.
|
|
weight: None or RaggedTensor which has the same shape of the input.
|
|
table: a table variable.
|
|
feature: a feature config.
|
|
|
|
Returns:
|
|
Embedding lookup result.
|
|
|
|
Raises:
|
|
ValueError: if input ragged tensor is not rank 2 or output shape set in
|
|
the feature config doesn't match with the first dim size of the input.
|
|
"""
|
|
if inp.shape.rank != 2:
|
|
raise ValueError(
|
|
"Only rank 2 ragged tensor is supported, but got rank {}".format(
|
|
inp.shape.rank))
|
|
batch_size = inp.shape[0]
|
|
|
|
# This computation needs to placed outside of tpu as the size of the row
|
|
# splits and values can change for different batch which can cause
|
|
# the program to re-compile.
|
|
def ragged_to_dense_outside_compilation(inp, weight, batch_size, feature):
|
|
if weight is None:
|
|
weight = ragged_tensor.RaggedTensor.from_row_splits(
|
|
array_ops.ones_like(inp.values, dtype=dtypes.float32),
|
|
inp.row_splits)
|
|
if not feature.output_shape and feature.max_sequence_length > 0:
|
|
inp = inp.to_tensor(shape=(batch_size, feature.max_sequence_length))
|
|
# Ignore weight if it is a sequence feature.
|
|
weight = array_ops.ones_like(inp, dtype=dtypes.float32)
|
|
elif feature.output_shape:
|
|
# Eagerly run the following op as the result as to be a number in
|
|
# order to use it as part of the output shape.
|
|
with ops.init_scope():
|
|
output_batch_size = math_ops.reduce_prod(feature.output_shape).numpy()
|
|
# If the output batch size matches the data batch size, treat it as
|
|
# normal ragged input.
|
|
if output_batch_size == batch_size:
|
|
inp, weight = inp.to_tensor(), weight.to_tensor()
|
|
# If the data batch size is a factor of the output batch size, the
|
|
# divide result will be the sequence length. Ignore the weights and
|
|
# combiner.
|
|
elif (
|
|
output_batch_size > batch_size
|
|
and output_batch_size % batch_size == 0
|
|
):
|
|
# Pad or truncate in the sequence dimension
|
|
seq_length = output_batch_size // batch_size
|
|
inp = inp.to_tensor(shape=(batch_size, seq_length))
|
|
# Ignore weight if it is a sequence feature.
|
|
weight = array_ops.ones_like(inp, dtype=dtypes.float32)
|
|
else:
|
|
raise ValueError(
|
|
"Output shape set in the FeatureConfig should be the factor of "
|
|
"the input data batch size. But instead got output shape {}, "
|
|
"input data batch size {}".format(feature.output_shape,
|
|
batch_size))
|
|
else:
|
|
inp, weight = inp.to_tensor(), weight.to_tensor()
|
|
return inp, weight
|
|
|
|
inp, weight = tpu_replication.outside_compilation(
|
|
ragged_to_dense_outside_compilation,
|
|
inp=inp,
|
|
weight=weight,
|
|
batch_size=batch_size,
|
|
feature=feature)
|
|
|
|
embeddings = embedding_ops.embedding_lookup_v2(table, inp)
|
|
weight = array_ops.expand_dims(weight, -1)
|
|
embeddings *= weight
|
|
|
|
if feature.output_shape:
|
|
with ops.init_scope():
|
|
output_batch_size = math_ops.reduce_prod(feature.output_shape).numpy()
|
|
if output_batch_size == batch_size:
|
|
embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
|
|
feature.table.combiner)
|
|
embeddings = array_ops.reshape(
|
|
embeddings, shape=feature.output_shape + [feature.table.dim])
|
|
else:
|
|
if feature.max_sequence_length == 0:
|
|
embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
|
|
feature.table.combiner)
|
|
return embeddings
|