def test_cpu_multiple_creation(self):
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends', max_sequence_length=2), )
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
embedding_one = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
embedding_two = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
# Both of the tpu embedding tables should be able to build on cpu.
embedding_one.build()
embedding_two.build()
def test_cpu_high_dimensional_lookup_ragged(self):
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends', output_shape=[2, 2]), )
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
features = self._get_ragged_tensors()[2:3]
result = mid_level(features, weights=None)
self.assertAllClose(result[0].shape, (2, 2, 2))
def test_cpu_no_optimizer(self):
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_video, name='watched', max_sequence_length=2), )
mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=None)
# Build the layer manually to create the variables. Normally calling enqueue
# would do this.
mid_level.build()
self.assertEqual(
list(mid_level._variables[self.table_video.name].keys()),
['parameters'])
def test_cpu_high_dimensional_sequence_lookup_ragged(self):
# Prod of output shape is a factor of the data batch size.
# The divide result will be the sequence length.
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends', output_shape=[2, 4]), )
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
features = self._get_ragged_tensors()[2:3]
result = mid_level(features, weights=None)
self.assertAllClose(result[0].shape, (2, 4, 2))
def test_cpu_high_dimensional_invalid_lookup_ragged(self):
# Prod of output shape is not a factor of the data batch size.
# An error will be raised in this case.
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends', output_shape=[3]), )
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
features = self._get_ragged_tensors()[2:3]
with self.assertRaisesRegex(
ValueError,
'Output shape set in the FeatureConfig should be the factor'):
mid_level(features, weights=None)
def test_cpu_sequence_lookup_sparse(self):
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends', max_sequence_length=2), )
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=feature_config, optimizer=optimizer)
features = self._get_sparse_tensors()[2:3]
result = mid_level(features, weights=None)
golden = self._numpy_sequence_lookup(
mid_level.embedding_tables[self.table_user].numpy(),
features[0].indices.numpy(), features[0].values.numpy(),
self.data_batch_size, feature_config[0].max_sequence_length,
self.table_user.dim)
self.assertAllClose(result[0], golden)
def test_check_checkpoint_variable_names_are_same_on_cpu_and_tpu(
self, optimizer):
# Reinitialize the TPU so that we can re-initialize the embeddings with the
# given optimizer.
if optimizer != tpu_embedding_v2_utils.SGD:
self.skip_if_oss()
strategy = self._get_strategy()
with strategy.scope():
first_mid_level_contents = np.ones((4, 4))
first_mid_level_optimizer = optimizer(learning_rate=0.1)
initializer = init_ops_v2.Constant(first_mid_level_contents)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=4,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
feature_config, first_mid_level_optimizer)
first_mid_level.build()
cpu_mid_level_optimizer = optimizer(learning_rate=0.1)
cpu_mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config, cpu_mid_level_optimizer)
cpu_mid_level.build()
tpu_checkpoint = util.Checkpoint(model=first_mid_level)
tpu_checkpoint.save(self._get_tmpdir('save-tpu', 'save'))
tpu_variables = checkpoint_utils.list_variables(
self._get_tmpdir('save-tpu'))
cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
cpu_checkpoint.save(self._get_tmpdir('save-cpu', 'save'))
cpu_variables = checkpoint_utils.list_variables(
self._get_tmpdir('save-cpu'))
self.assertAllEqual(tpu_variables, cpu_variables)
def test_model_export_cpu(self):
strategy = self._get_strategy()
with strategy.scope():
first_mid_level_contents = np.ones((4, 4))
first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
initializer = init_ops_v2.Constant(first_mid_level_contents)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=4,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
feature_config, first_mid_level_optimizer)
first_mid_level.build()
cpu_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
cpu_mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config, cpu_mid_level_optimizer)
cpu_mid_level.build()
tpu_checkpoint = util.Checkpoint(model=first_mid_level)
tpu_checkpoint.save(self._get_tmpdir('export_cpu', 'save'))
# We restore the checkpoint of our tpu mid level onto our cpu mid level.
cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
cpu_checkpoint.restore(self._get_tmpdir('export_cpu', 'save-1'))
@def_function.function
def serve_tensors(features):
features = tpu_embedding_for_serving.cpu_embedding_lookup(
features, None, cpu_mid_level.embedding_tables,
cpu_mid_level._feature_config)
return features[0]
signatures = {
'serving_default':
serve_tensors.get_concrete_function((tensor_spec.TensorSpec(
shape=(2,), dtype=dtypes.int32, name='feature'),))
}
save.save(
cpu_mid_level,
export_dir=self._get_tmpdir('export_cpu', 'exported_model'),
signatures=signatures)
imported = load.load(self._get_tmpdir('export_cpu', 'exported_model'))
predict_fn = imported.signatures['serving_default']
input_feature_value = np.array([1, 0])
input_batch = (constant_op.constant(
input_feature_value, dtype=dtypes.int32),)
prediction = predict_fn(*input_batch)['output_0']
self.assertAllClose(prediction.numpy(),
first_mid_level_contents[input_feature_value])
def _create_mid_level(self):
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
return tpu_embedding_for_serving.TPUEmbeddingForServing(
feature_config=self.feature_config, optimizer=optimizer)