def test_shared_sequence_non_sequence_into_input_layer(self):
non_seq = fc.categorical_column_with_identity('non_seq',
num_buckets=10)
seq = sfc.sequence_categorical_column_with_identity('seq',
num_buckets=10)
shared_non_seq, shared_seq = fc.shared_embedding_columns_v2(
[non_seq, seq],
dimension=4,
combiner='sum',
initializer=init_ops_v2.Ones(),
shared_embedding_collection_name='shared')
seq = sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
values=[0, 1, 2],
dense_shape=[2, 2])
non_seq = sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
values=[0, 1, 2],
dense_shape=[2, 2])
features = {'seq': seq, 'non_seq': non_seq}
# Tile the context features across the sequence features
seq_input, seq_length = ksfc.SequenceFeatures([shared_seq])(features)
non_seq_input = dense_features.DenseFeatures([shared_non_seq
])(features)
with self.cached_session() as sess:
sess.run(variables.global_variables_initializer())
output_seq, output_seq_length, output_non_seq = sess.run(
[seq_input, seq_length, non_seq_input])
self.assertAllEqual(
output_seq,
[[[1, 1, 1, 1], [1, 1, 1, 1]], [[1, 1, 1, 1], [0, 0, 0, 0]]])
self.assertAllEqual(output_seq_length, [2, 1])
self.assertAllEqual(output_non_seq, [[2, 2, 2, 2], [1, 1, 1, 1]])
def test_checkpoint_restore_before_variable_creation(self):
class TestModule(module.Module):
def __init__(self, initializer, rows):
self._initializer = initializer
self._rows = rows
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=self._rows,
dim=4,
initializer=self._initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
optimizer = tpu_embedding_v2_utils.SGD()
self.tpu_embedding = tpu_embedding_v2.TPUEmbedding(
feature_config, optimizer)
def create_embedding(self):
# We aren't training so batch_size here doesn't matter.
self.tpu_embedding.build(64)
strategy = self._get_strategy()
with strategy.scope():
module1 = TestModule(init_ops_v2.Ones(),
strategy.num_replicas_in_sync * 2)
module1.create_embedding()
checkpoint = util.Checkpoint(test_module=module1)
checkpoint.save(self._get_tmpdir('restore_before_create', 'save'))
# Reinitialize the tpu
strategy = self._get_strategy()
with strategy.scope():
module2 = TestModule(init_ops_v2.Zeros(),
strategy.num_replicas_in_sync * 2)
checkpoint = util.Checkpoint(test_module=module2)
checkpoint.restore(self._get_tmpdir('restore_before_create', 'save-1'))
with strategy.scope():
module2.create_embedding()
def get_values(mid):
return mid._variables['table']['parameters'].variables[0].numpy()
self.assertAllClose(
np.ones((strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
# Fetch the values from the TPU to check that they are the same.
module2.tpu_embedding._retrieve_variables()
self.assertAllClose(
np.ones((strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
def test_checkpoint_restore_before_variable_creation(self):
# This test works right now because we only have one TPU host in the unit
# environment. Initializing from checkpoint does not understand how to
# pass the sharding info to the restore op right now.
class TestModule(module.Module):
def __init__(self, initializer, rows):
self._initializer = initializer
self._rows = rows
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=self._rows, dim=4, initializer=self._initializer,
combiner='sum', name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
optimizer = tpu_embedding_v2_utils.SGD()
self.tpu_embedding = tpu_embedding_v2.TPUEmbedding(
feature_config, optimizer)
def create_embedding(self):
# We aren't training so batch_size here doesn't matter.
self.tpu_embedding.build(64)
# We need to clear the any already loaded config provided by setUp method.
tpu_strategy_util.initialize_tpu_system(self.resolver)
with self.strategy.scope():
module1 = TestModule(init_ops_v2.Ones(),
self.strategy.num_replicas_in_sync * 2)
module1.create_embedding()
checkpoint = util.Checkpoint(test_module=module1)
checkpoint.save(_get_tmpdir('restore_before_create', 'save'))
tpu_strategy_util.initialize_tpu_system(self.resolver)
with self.strategy.scope():
module2 = TestModule(init_ops_v2.Zeros(),
self.strategy.num_replicas_in_sync * 2)
checkpoint = util.Checkpoint(test_module=module2)
checkpoint.restore(_get_tmpdir('restore_before_create', 'save-1'))
with self.strategy.scope():
module2.create_embedding()
def get_values(mid):
return mid._variables['table']['parameters'].variables[0].numpy()
self.assertAllClose(np.ones((self.strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
# Fetch the values from the TPU to check that they are the same.
module2.tpu_embedding._retrieve_variables()
self.assertAllClose(np.ones((self.strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
def test_checkpoint_restore_before_variable_creation(self):
class TestModule(module.Module):
def __init__(self, initializer, rows):
self._initializer = initializer
self._rows = rows
def create_embedding(self):
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=self._rows, dim=4, initializer=self._initializer,
combiner='sum', name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
optimizer = tpu_embedding_v2_utils.SGD()
self.tpu_embedding = tpu_embedding_v2.TPUEmbedding(
feature_config, self._rows, optimizer)
# We need to clear the already loaded config provided by setUp method.
tpu_strategy_util.initialize_tpu_system(self.resolver)
with self.strategy.scope():
module1 = TestModule(init_ops_v2.Ones(),
self.strategy.num_replicas_in_sync * 2)
module1.create_embedding()
checkpoint = util.Checkpoint(test_module=module1)
checkpoint.save(_get_tmpdir('restore_before_create', 'save'))
tpu_strategy_util.initialize_tpu_system(self.resolver)
with self.strategy.scope():
module2 = TestModule(init_ops_v2.Zeros(),
self.strategy.num_replicas_in_sync * 2)
checkpoint = util.Checkpoint(test_module=module2)
checkpoint.restore(_get_tmpdir('restore_before_create', 'save-1'))
with self.strategy.scope():
module2.create_embedding()
def get_values(mid):
return mid._variables['table']['parameters'].variables[0].numpy()
self.assertAllClose(np.ones((self.strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
# Fetch the values from the TPU to check that they are the same.
module2.tpu_embedding._retrieve_variables()
self.assertAllClose(np.ones((self.strategy.num_replicas_in_sync * 2, 4)),
get_values(module2.tpu_embedding))
def testOnesPartition(self): init = init_ops_v2.Ones() self._partition_test(init)
def testOnes(self):
self._range_test(init_ops_v2.Ones(), shape=(4, 5),
target_mean=1., target_max=1.)
def build(self):
self.w = self._add_variable_with_custom_getter(
"w",
shape=(4,),
initializer=init_ops_v2.Ones(),
getter=make_variable)
