def testPartitionToOne(self):
# For small variables there is only one partition.
variable_partitioner = partitioned_variables.min_max_variable_partitioner(
max_partitions=2, min_slice_size=64 << 20)
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver, variable_partitioner)
with strategy.scope():
initializer = init_ops_v2.Constant([0] * 10)
v1 = variables.Variable(initial_value=lambda: initializer(
shape=(10, ), dtype=dtypes.int64),
shape=(10, ),
dtype=dtypes.int64)
v2 = variables.Variable([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
self.assertIsInstance(v1, variables.Variable)
self.assertNotIsInstance(v1, sharded_variable.ShardedVariable)
self.assertRegex(v1.device, "/job:ps/replica:0/task:0")
self.assertAllEqual(v1.read_value().numpy(), [0] * 10)
self.assertIsInstance(v2, variables.Variable)
self.assertNotIsInstance(v2, sharded_variable.ShardedVariable)
self.assertRegex(v2.device, "/job:ps/replica:0/task:1")
self.assertAllEqual(v2.read_value().numpy(),
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
def testPartitionWhenLackOfInfo(self):
strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
self.cluster_resolver,
partitioned_variables.fixed_size_partitioner(2))
with strategy.scope():
initializer = init_ops_v2.Constant([0, 1, 2, 3])
# Shape is not explicitly specified.
v1 = variables.Variable(initial_value=lambda: initializer(
shape=(4, ), dtype=dtypes.int64),
dtype=dtypes.int64)
# Dtype is not explicitly specified.
v2 = variables.Variable(initial_value=lambda: initializer(
shape=(4, ), dtype=dtypes.int64),
shape=(4, ))
# Neither shape nor dtype is explicitly specified.
v3 = variables.Variable(initial_value=lambda: initializer(
shape=(4, ), dtype=dtypes.int64))
for v in [v1, v2, v3]:
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 2)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertRegex(v.variables[1].device, "/job:ps/replica:0/task:1")
self.assertAllEqual(v.variables[0].read_value().numpy(), [0, 1])
self.assertAllEqual(v.variables[1].read_value().numpy(), [2, 3])
def build_mid_level(self, embedding_values, optimizer,
initialize_tpu_embedding=True):
"""Creates an embedding api object initialized to embedding_values."""
initializer = init_ops_v2.Constant(embedding_values)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=self.num_rows, dim=4, initializer=initializer,
combiner='sum', name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, optimizer)
# We want to create a second object (with its own variables) but not
# initialize the TPU.
if not initialize_tpu_embedding:
saved_fn = tpu.initialize_system_for_tpu_embedding
tpu.initialize_system_for_tpu_embedding = lambda x: None
# batch_size here does not matter as we aren't training in any of these
# tests.
mid_level.build(64)
if not initialize_tpu_embedding:
tpu.initialize_system_for_tpu_embedding = saved_fn
return mid_level
def testConstantPartition(self):
init = init_ops_v2.Constant([1, 2, 3, 4])
with self.assertRaisesWithLiteralMatch(
ValueError,
r"Constant initializer doesn't support partition-related arguments"
):
init((4, 2), dtype=dtypes.float32, partition_shape=(2, 2))
def testConstantInt(self):
self._range_test(
init_ops_v2.Constant(2),
shape=(5, 6, 4),
target_mean=2,
target_max=2,
target_min=2)
def test_checkpoint_save_and_restore(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()
first_checkpoint = util.Checkpoint(model=first_mid_level)
first_checkpoint.save(self._get_tmpdir('restore', 'save'))
with strategy.scope():
second_mid_level_contents = np.ones((4, 4)) * 2
second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
initializer = init_ops_v2.Constant(second_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'),)
second_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
feature_config, second_mid_level_optimizer)
second_mid_level.build()
# We restore the checkpoint of our first model into our second model.
second_checkpoint = util.Checkpoint(model=second_mid_level)
second_checkpoint.restore(self._get_tmpdir('restore', 'save-1'))
self.assertAllClose(
first_mid_level_contents,
second_mid_level._variables['table']['parameters'].numpy(),
msg='Second mid level api should have restored the first model values.')
def testInvalidArgument(self):
with self.assertRaisesRegex(ValueError, "initial_value"):
with self.client.experimental_variable_partitioning_scope():
variables.Variable(initial_value=[0, 1, 2], shape=(3, ))
with self.assertRaisesRegex(ValueError, "shape"):
with self.client.experimental_variable_partitioning_scope():
initializer = init_ops_v2.Constant([0, 1, 2])
variables.Variable(initial_value=lambda: initializer(
shape=(3, ), dtype=dtypes.int64),
dtype=dtypes.int64)
def testSurplusPS(self):
with self.client.strategy.scope():
with self.client.experimental_variable_partitioning_scope():
initializer = init_ops_v2.Constant([0])
v = variables.Variable(initial_value=lambda: initializer(
shape=(1, ), dtype=dtypes.int64),
shape=(1, ),
dtype=dtypes.int64)
self.assertIsInstance(v, sharded_variable.ShardedVariable)
self.assertLen(v.variables, 1)
self.assertRegex(v.variables[0].device, "/job:ps/replica:0/task:0")
self.assertAllEqual(v.variables[0].read_value().numpy(), [0])
def build_mid_level(self, embedding_values, optimizer,
initialize_tpu_embedding=True):
"""Creates an embedding api object initialized to embedding_values."""
initializer = init_ops_v2.Constant(embedding_values)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=self.num_rows, dim=4, initializer=initializer,
combiner='sum', name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
# batch_size here does not matter as we aren't training in any of these
# tests.
return tpu_embedding_v2.TPUEmbedding(
feature_config, 64, optimizer,
initialize_tpu_embedding=initialize_tpu_embedding)
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()
num_rows = strategy.num_replicas_in_sync
with strategy.scope():
first_mid_level_contents = np.ones((num_rows, 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=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, first_mid_level_optimizer)
first_mid_level.build(64)
cpu_mid_level_optimizer = optimizer(learning_rate=0.1)
cpu_mid_level = tpu_embedding_v2.TPUEmbedding(feature_config,
cpu_mid_level_optimizer)
cpu_mid_level.build(64)
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 _slot_initializers(self) -> List[init_ops_v2.Initializer]:
return [
init_ops_v2.Constant(self.initial_accumulator_value),
init_ops_v2.Constant()
]
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 setUp(self):
super(TPUEmbeddingBaseTest, self).setUp()
self.embedding_values = np.array(list(range(32)), dtype=np.float64)
self.initializer = init_ops_v2.Constant(self.embedding_values)
# Embedding for video initialized to
# 0 1 2 3
# 4 5 6 7
# ...
self.table_video = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=8,
dim=4,
initializer=self.initializer,
combiner='sum',
name='video')
# Embedding for user initialized to
# 0 1
# 2 3
# 4 5
# 6 7
# ...
self.table_user = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=16,
dim=2,
initializer=self.initializer,
combiner='mean',
name='user')
self.feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_video, name='watched'),
tpu_embedding_v2_utils.FeatureConfig(
table=self.table_video, name='favorited'),
tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends'))
self.batch_size = 2
self.data_batch_size = 4
# One (global) batch of inputs
# sparse tensor for watched:
# row 0: 0
# row 1: 0, 1
# row 2: 0, 1
# row 3: 1
self.feature_watched_indices = [[0, 0], [1, 0], [1, 1], [2, 0], [2, 1],
[3, 0]]
self.feature_watched_values = [0, 0, 1, 0, 1, 1]
self.feature_watched_row_lengths = [1, 2, 2, 1]
# sparse tensor for favorited:
# row 0: 0, 1
# row 1: 1
# row 2: 0
# row 3: 0, 1
self.feature_favorited_indices = [[0, 0], [0, 1], [1, 0], [2, 0],
[3, 0], [3, 1]]
self.feature_favorited_values = [0, 1, 1, 0, 0, 1]
self.feature_favorited_row_lengths = [2, 1, 1, 2]
# sparse tensor for friends:
# row 0: 3
# row 1: 0, 1, 2
# row 2: 3
# row 3: 0, 1, 2
self.feature_friends_indices = [[0, 0], [1, 0], [1, 1], [1, 2], [2, 0],
[3, 0], [3, 1], [3, 2]]
self.feature_friends_values = [3, 0, 1, 2, 3, 0, 1, 2]
self.feature_friends_row_lengths = [1, 3, 1, 3]
self.resolver = None
# Basically we are expand the dims of the old feature by 1 and repeat
# batch size times for the first dimension.
def create_hight_dimensional_indices(indices):
indices = np.array(indices, dtype=np.int32)
batch_size_index = np.repeat(np.arange(self.data_batch_size),
len(indices)).reshape(-1, 1)
repeated_indices = np.tile(indices, (self.data_batch_size, 1))
return np.concatenate([batch_size_index, repeated_indices], axis=1)
# Create high dimensional features with shape(4, 4, 2)
self.feature_watched_indices_high_dimensional = create_hight_dimensional_indices(
self.feature_watched_indices)
self.feature_watched_values_high_dimensional = self.feature_watched_values * self.data_batch_size
self.feature_watched_row_lengths_high_dimensional = self.feature_watched_row_lengths * self.data_batch_size
# Create high dimensional features with shape(4, 4, 2)
self.feature_favorited_indices_high_dimensional = create_hight_dimensional_indices(
self.feature_favorited_indices)
self.feature_favorited_values_high_dimensional = self.feature_favorited_values * self.data_batch_size
self.feature_favorited_row_lengths_high_dimensional = self.feature_favorited_row_lengths * self.data_batch_size
# Create high dimensional features with shape(4, 4, 3)
self.feature_friends_indices_high_dimensional = create_hight_dimensional_indices(
self.feature_friends_indices)
self.feature_friends_values_high_dimensional = self.feature_friends_values * self.data_batch_size
self.feature_friends_row_lengths_high_dimensional = self.feature_friends_row_lengths * self.data_batch_size
def _slot_initializers(self):
return [init_ops_v2.Constant(), init_ops_v2.Constant()]
def _slot_initializers(self):
return [init_ops_v2.Constant(self.initial_accumulator_value)]
def test_checkpoint_restore_loads(self):
strategy = self._get_strategy()
num_rows = strategy.num_replicas_in_sync
def get_values(mid):
return ops.convert_to_tensor(
mid._variables['table']['parameters'].variables[0])
with strategy.scope():
first_mid_level_contents = np.ones((num_rows, 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=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, first_mid_level_optimizer)
first_mid_level.build(64)
first_mid_level._load_variables()
first_checkpoint = util.Checkpoint(model=first_mid_level)
first_checkpoint.save(self._get_tmpdir('restore', 'save'))
tpu_strategy_util.initialize_tpu_system(self.resolver)
with strategy.scope():
second_mid_level_contents = np.ones((num_rows, 4)) * 2
second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
initializer = init_ops_v2.Constant(second_mid_level_contents)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
second_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, second_mid_level_optimizer)
second_mid_level.build(64)
second_mid_level._load_variables()
self.assertAllClose(
second_mid_level_contents,
get_values(second_mid_level),
msg='Second mid level api should contain its initial values.',
)
# We restore the checkpoint of our first model into our second model.
# This should load the first mid level API object onto the TPU.
second_checkpoint = util.Checkpoint(model=second_mid_level)
second_checkpoint.restore(self._get_tmpdir('restore', 'save-1'))
# Call retrieve here as a way to check what the TPU contains.
# Calling the retrieve ops directly might make for a cleaner separation of
# test and module, though.
second_mid_level._retrieve_variables()
self.assertAllClose(
first_mid_level_contents,
get_values(second_mid_level),
msg='Second mid level api should have retrieved the first model values.'
)
def setUp(self):
super(CPUEmbeddingTest, self).setUp()
self.embedding_values = np.array(list(range(32)), dtype=np.float64)
self.initializer = init_ops_v2.Constant(self.embedding_values)
# Embedding for video initialized to
# 0 1 2 3
# 4 5 6 7
# ...
self.table_video = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=8,
dim=4,
initializer=self.initializer,
combiner='sum',
name='video')
# Embedding for user initialized to
# 0 1
# 2 3
# 4 5
# 6 7
# ...
self.table_user = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=16,
dim=2,
initializer=self.initializer,
combiner='mean',
name='user')
self.feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=self.table_video, name='watched'),
tpu_embedding_v2_utils.FeatureConfig(
table=self.table_video, name='favorited'),
tpu_embedding_v2_utils.FeatureConfig(
table=self.table_user, name='friends'))
self.batch_size = 2
self.data_batch_size = 4
# One (global) batch of inputs
# sparse tensor for watched:
# row 0: 0
# row 1: 0, 1
# row 2: 0, 1
# row 3: 1
self.feature_watched_indices = [[0, 0], [1, 0], [1, 1], [2, 0], [2, 1],
[3, 0]]
self.feature_watched_values = [0, 0, 1, 0, 1, 1]
self.feature_watched_row_lengths = [1, 2, 2, 1]
# sparse tensor for favorited:
# row 0: 0, 1
# row 1: 1
# row 2: 0
# row 3: 0, 1
self.feature_favorited_indices = [[0, 0], [0, 1], [1, 0], [2, 0],
[3, 0], [3, 1]]
self.feature_favorited_values = [0, 1, 1, 0, 0, 1]
self.feature_favorited_row_lengths = [2, 1, 1, 2]
# sparse tensor for friends:
# row 0: 3
# row 1: 0, 1, 2
# row 2: 3
# row 3: 0, 1, 2
self.feature_friends_indices = [[0, 0], [1, 0], [1, 1], [1, 2], [2, 0],
[3, 0], [3, 1], [3, 2]]
self.feature_friends_values = [3, 0, 1, 2, 3, 0, 1, 2]
self.feature_friends_row_lengths = [1, 3, 1, 3]
def test_variable_learning_rate(self):
num_steps = 10
num_steps_float = float(num_steps)
starting_lr = 1.0
ending_lr = 0.5
strategy = self._get_strategy()
num_replicas = strategy.num_replicas_in_sync
# Create model with Keras.
with strategy.scope():
step_counter = tf_variables.Variable(0.0, dtypes.float32)
def lr_function():
return gen_math_ops.maximum(
ending_lr,
starting_lr + ((ending_lr - starting_lr) * step_counter) /
num_steps_float)
optimizer = tpu_embedding_v2_utils.SGD(learning_rate=lr_function)
table_config = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=num_replicas,
dim=4,
initializer=init_ops_v2.Constant(np.zeros((num_replicas, 4))),
combiner='sum', name='table')
mid_level_api = tpu_embedding_v2.TPUEmbedding(
feature_config={
'feature': tpu_embedding_v2_utils.FeatureConfig(
table=table_config, name='feature')},
optimizer=optimizer)
feature = {
'feature': constant_op.constant([0], shape=(1, 1), dtype=dtypes.int32)
}
def input_fn(ctx):
del ctx
return dataset_ops.DatasetV2.from_tensors(feature).repeat()
dist = strategy.distribute_datasets_from_function(
input_fn,
options=distribute_lib.InputOptions(experimental_fetch_to_device=False))
dist_iter = iter(dist)
@def_function.function
def test_fn():
def step():
with backprop.GradientTape() as tape:
activations = mid_level_api.dequeue()
tape.watch(activations)
result = math_ops.reduce_sum(activations['feature'])
loss = result / num_replicas
grads = tape.gradient(loss, activations)
mid_level_api.apply_gradients(grads)
return activations['feature']
mid_level_api.enqueue(next(dist_iter), training=True)
return strategy.run(step)
# Run model.
results = []
for _ in range(num_steps):
result = test_fn()
results.append(self._unpack(strategy, result))
step_counter.assign_add(1.0)
# Table is 2 elements wide, per-replica batch size of 1, with id 0.
# Loss for the gradient is the sum of the entries divided by the number of
# replicas. Thus the per replica gradient is 1/#of replicas for row 0 and no
# other updates. The reduced gradient is therefore 1.
# Learning rate schedule over num_steps steps:
# 1.0 0.95 0.9 0.85 0.8 ...
# Since use SGD and the gradient is one, the first row of the table is
# [0, 0] [-1.0, -1.0] [-1.95, -1.95] [-2.85, -2.85] ... (the negative
# partial sums of the above).
learning_rates = [starting_lr - (starting_lr - ending_lr) / num_steps * j
for j in range(num_steps)]
cumsum = [sum(learning_rates[0:j]) for j in range(num_steps)]
goldens = [[[-cumsum[i]] * table_config.dim] * num_replicas
for i in range(10)]
self.assertAllClose(results, goldens)
def test_checkpoint_save_retrieves(self):
strategy = self._get_strategy()
num_rows = strategy.num_replicas_in_sync
with strategy.scope():
first_mid_level_contents = np.ones((num_rows, 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=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, first_mid_level_optimizer)
first_mid_level.build(64)
# Ensure that the variables from the first model are loaded.
first_mid_level._load_variables()
self.assertAllClose(
first_mid_level_contents,
self.make_checkpoint_and_get_embedding('before_load', first_mid_level,
num_rows),
msg='Checkpoint should contain values from the first api object.')
# Reinitialize the tpu.
tpu_strategy_util.initialize_tpu_system(self.resolver)
with strategy.scope():
second_mid_level_contents = np.ones((num_rows, 4)) * 2
second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
initializer = init_ops_v2.Constant(second_mid_level_contents)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
second_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, second_mid_level_optimizer)
second_mid_level.build(64)
second_mid_level._load_variables()
# When we load the variables from the second mid level API object to the TPU
# we expect that checkpointing the first mid level API object will now
# retrieve the values from the TPU which are now different from the current
# variables in the first mid level.
self.assertAllClose(
second_mid_level_contents,
self.make_checkpoint_and_get_embedding('after_load', first_mid_level,
num_rows),
msg='Checkpoint should contain values from the second api object.')
def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
return [init_ops_v2.Constant(), init_ops_v2.Constant()]
parser.add_argument("--features", type=int, default=10)
parser.add_argument("--em", type=int, default=4)
parser.add_argument("--nnz", type=int, default=2)
parser.add_argument("--batch", type=int, default=4)
parser.add_argument("--steps", type=int, default=1)
parser.add_argument("--warmups", type=int, default=0)
parser.add_argument("--randomseed", type=int, default=0)
parser.add_argument("--testtpu", type=int, default=0)
parser.add_argument("--verify", type=int, default=0)
parser.add_argument("--times", type=int, default=1)
args = parser.parse_args()
# here
embedding_values = np.array(list(range(40)), dtype=np.float64)
initializer = init_ops_v2.Constant(embedding_values)
table_test = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=args.features,
dim=args.em,
# initializer=initializer,
initializer=None,
combiner='sum',
name='test')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(table=table_test,
name='watched'))
batch = args.batch
nnz = args.nnz
features = args.features
