def test_save_graph_def(self):
root = tracking.AutoTrackable()
v1 = variables_lib.Variable([3.])
v2 = variables_lib.Variable([2.])
root.v = sharded_variable.ShardedVariable([v1, v2])
root.train = def_function.function(
lambda x: embedding_ops.embedding_lookup_v2(root.v.variables, x))
# TODO(b/144057383): Remove the necessity of root.serve once saving context
# is made to tf.function cache.
root.serve = def_function.function(
lambda x: embedding_ops.embedding_lookup_v2(
root.v.variables[0], x),
input_signature=[
tensor_spec.TensorSpec([2], dtypes.int32, name='x')
])
# Trace and use root.train
self.assertAllEqual([3., 2.], root.train([0, 1]).numpy())
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
save.save(root, save_dir, root.serve)
self.assertAllEqual([3., 2.],
_load_and_run(save_dir, {'x': [0, 1]})['output_0'])
# Continue using root.train for training
self.assertAllEqual([3., 2.], root.train([0, 1]).numpy())
def _embedding_call(self, og_call, layer, inputs, *args, **kwargs):
embeddings = layer.embeddings
# NOTE: We need to put this in two if statements as autograph is dumb.
if embeddings.ref() not in self.ref_to_logvar:
# This can happen for non-mergeable variables, so we skip them.
return og_call(layer, inputs, *args, **kwargs)
if not self.sampling:
return og_call(layer, inputs, *args, **kwargs)
embeddings_logvar = self.ref_to_logvar[embeddings.ref()]
stddev = tf.exp(0.5 * embeddings_logvar)
batch_size = tf.shape(inputs)[0]
# TODO: Find a better way of doing all of this.
output = tf.zeros_like(embedding_ops.embedding_lookup_v2(embeddings, inputs))
for i in tf.range(batch_size):
eps = tf.random.normal(tf.shape(embeddings))
perturbed_embeddings = embeddings + stddev * eps
mask = tf.one_hot(i, depth=batch_size)
output_i = embedding_ops.embedding_lookup_v2(perturbed_embeddings, inputs)
for _ in range(len(output_i.shape) - 1):
mask = mask[..., None]
output += output_i * mask
return output
def call(self, inputs):
dtype = K.dtype(inputs)
if dtype != 'int32' and dtype != 'int64':
inputs = math_ops.cast(inputs, 'int32')
if isinstance(self.embeddings, sharded_variable.ShardedVariable):
out = embedding_ops.embedding_lookup_v2(self.embeddings.variables, inputs)
else:
out = embedding_ops.embedding_lookup_v2(self.embeddings, inputs)
return out
def call(self, inputs):
dtype = K.dtype(inputs)
if dtype != 'int32' and dtype != 'int64':
inputs = math_ops.cast(inputs, 'int32')
if isinstance(self.embeddings, sharded_variable.ShardedVariable):
out = embedding_ops.embedding_lookup_v2(self.embeddings.variables,
inputs)
else:
out = embedding_ops.embedding_lookup_v2(self.embeddings, inputs)
if self._dtype_policy.should_cast_variables:
# Instead of casting the variable as in most layers, cast the output, as
# this is mathematically equivalent but is faster.
out = math_ops.cast(out, self._dtype_policy.compute_dtype)
return out
def _call():
with gradients.GradientTape() as tape:
y = embedding_ops.embedding_lookup_v2(x, [0])
loss = math_ops.reduce_sum(y)
grads = tape.gradient(loss, x)
self.assertAllEqual(grads.shape, [3, 3])
return ops.convert_to_tensor(grads)
def _handle_read(new_graph_item, var_op, partitioned_var):
partitioned_var_tensor = partitioned_var.as_tensor()
for op in get_consumers(var_op):
op = new_graph_item.graph.get_operation_by_name(
ops.prepend_name_scope(op.name, AUTODIST_TO_DELETE_SCOPE)
)
if op.type == "ResourceGather":
# Only Resource Variable needs to be taken care of
# because ResourceGather consumes resource tensor rather than the tensor of read_var_op
# Question: Is there any case where the op.type == "ResourceGather"
# but we can't use embedding_lookup_v2 to reconstruct the op consuming a partitioned resource
# The second input to a ResourceGather op is always the indices per the opdef
emb_lookup = embedding_ops.embedding_lookup_v2(partitioned_var, ids=op.inputs[1])
update_consumers(get_consumers(op), op.outputs[0], emb_lookup)
if is_read_var_op(op, version=1):
# Without our modification, Reference Vars in TF have a read op associated with them.
# TF can sometimes look for this and expect it to exist (e.g. in graph.as_graph_element)
# so we add one back to avoid errors.
# read_out is already the output tensor of the generated identity op
read_out = array_ops.identity(partitioned_var_tensor,
name=ops.prepend_name_scope("read", var_op.name))
update_consumers(get_consumers(op), op.outputs[0], read_out)
elif is_read_var_op(op, version=2):
read_out = array_ops.identity(partitioned_var_tensor,
name=ops.prepend_name_scope("Read/ReadVariableOp", var_op.name))
update_consumers(get_consumers(op), op.outputs[0], read_out)
def avs_embedding(self, y):
if 'word_id' in y:
ids = tf.argmax(y['word_id'], axis=-1)
y_vec = embedding_ops.embedding_lookup_v2(self.embeddings, ids)
else:
y_word = y2action_word(y)
y_vec = tf.numpy_function(self.word2word_vec, [y_word], tf.float32)
return y_vec
def call(self, inputs):
dtype = backend.dtype(inputs)
if dtype != 'int32' and dtype != 'int64':
inputs = math_ops.cast(inputs, 'int32')
out = embedding_ops.embedding_lookup_v2(self.embeddings, inputs)
if self._dtype_policy.compute_dtype != self._dtype_policy.variable_dtype:
# Instead of casting the variable as in most layers, cast the output, as
# this is mathematically equivalent but is faster.
out = math_ops.cast(out, self._dtype_policy.compute_dtype)
return out
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) -> ops.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.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 lookup(): ids = constant_op.constant([0, 3, 4]) return embedding_ops.embedding_lookup_v2(sv, ids)
def _call():
return embedding_ops.embedding_lookup_v2(x, [0])
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) -> ops.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.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
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, ops.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, ops.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_ragged_tensor(
inp: ragged_tensor.RaggedTensor,
weight: Optional[ragged_tensor.RaggedTensor],
table: tf_variables.Variable,
feature: tpu_embedding_v2_utils.FeatureConfig) -> ops.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]
if feature.output_shape:
output_batch_size = math_ops.reduce_prod(feature.output_shape)
# If the output batch size matches the data batch size, treat it as
# normal ragged input.
if output_batch_size == batch_size:
ragged_output = _ragged_embedding_lookup_with_reduce(
table, inp, weight, feature.table.combiner)
ragged_output = array_ops.reshape(ragged_output,
shape=feature.output_shape +
[feature.table.dim])
# 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:
ragged_output = embedding_ops.embedding_lookup_v2(table, inp)
# Pad or truncate in the sequence dimension
ragged_output = ragged_output.to_tensor(shape=[
batch_size, output_batch_size // batch_size, feature.table.dim
])
# Reshape to desire output shape.
ragged_output = array_ops.reshape(
ragged_output, feature.output_shape + [feature.table.dim])
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:
if feature.max_sequence_length > 0:
output_shape = [
batch_size, feature.max_sequence_length, feature.table.dim
]
ragged_lookup = embedding_ops.embedding_lookup_v2(table, inp)
# Unlike scatter_nd, RaggedTensor.to_tensor truncates to the given
# shape.
ragged_output = ragged_lookup.to_tensor(shape=output_shape)
else:
ragged_output = _ragged_embedding_lookup_with_reduce(
table, inp, weight, feature.table.combiner)
return ragged_output
def cpu_embedding_lookup(
inputs: Any,
weights: Optional[Any],
tables: Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable],
feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable] # pylint:disable=g-bare-generic
) -> Any:
"""Apply standard lookup ops with `tf.tpu.experimental.embedding` configs.
This function is a utility which allows using the
`tf.tpu.experimental.embedding` config objects with standard lookup functions.
This can be used when exporting a model which uses
`tf.tpu.experimental.embedding.TPUEmbedding` for serving on CPU. In particular
`tf.tpu.experimental.embedding.TPUEmbedding` only supports lookups on TPUs and
should not be part of your serving graph.
Note that TPU specific options (such as `max_sequence_length`) in the
configuration objects will be ignored.
In the following example we take a trained model (see the documentation for
`tf.tpu.experimental.embedding.TPUEmbedding` for the context) and create a
saved model with a serving function that will perform the embedding lookup and
pass the results to your model:
```python
model = model_fn(...)
embedding = tf.tpu.experimental.embedding.TPUEmbedding(
feature_config=feature_config,
batch_size=1024,
optimizer=tf.tpu.experimental.embedding.SGD(0.1))
checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
checkpoint.restore(...)
@tf.function(input_signature=[{'feature_one': tf.TensorSpec(...),
'feature_two': tf.TensorSpec(...),
'feature_three': tf.TensorSpec(...)}])
def serve_tensors(embedding_features):
embedded_features = tf.tpu.experimental.embedding.serving_embedding_lookup(
embedding_features, None, embedding.embedding_tables,
feature_config)
return model(embedded_features)
model.embedding_api = embedding
tf.saved_model.save(model,
export_dir=...,
signatures={'serving_default': serve_tensors})
```
NOTE: It's important to assign the embedding API object to a member of your
model as `tf.saved_model.save` only supports saving variables as one
`Trackable` object. Since the model's weights are in `model` and the
embedding table are managed by `embedding`, we assign `embedding` to an
attribute of `model` so that tf.saved_model.save can find the embedding
variables.
NOTE: The same `serve_tensors` function and `tf.saved_model.save` call will
work directly from training.
Args:
inputs: 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.
tables: a dict of mapping TableConfig objects to Variables.
feature_config: a nested structure of FeatureConfig objects with the same
structure as inputs.
Returns:
A nested structure of Tensors with the same structure as inputs.
"""
nest.assert_same_structure(inputs, feature_config)
flat_inputs = nest.flatten(inputs)
flat_weights = [None] * len(flat_inputs)
if weights is not None:
nest.assert_same_structure(inputs, weights)
flat_weights = nest.flatten(weights)
flat_features = nest.flatten_with_joined_string_paths(feature_config)
outputs = []
for inp, weight, (path, feature) in zip(flat_inputs, flat_weights,
flat_features):
table = tables[feature.table]
if weight is not None:
if isinstance(inp, ops.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, ops.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(
_embedding_lookup_for_sparse_tensor(inp, weight, table,
feature))
elif isinstance(inp, ragged_tensor.RaggedTensor):
outputs.append(
_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(feature_config, outputs)