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)