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 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 test_table_config_repr(self): table = tpu_embedding_v2_utils.TableConfig( vocabulary_size=2, dim=4, initializer=None, combiner='sum', name='table') self.assertEqual( repr(table), 'TableConfig(vocabulary_size=2, dim=4, initializer=None, ' 'optimizer=None, combiner=\'sum\', name=\'table\')')
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)
def test_tables_with_same_name(self): with self.assertRaisesRegex( ValueError, 'Multiple tables with name table found.'): with self._get_strategy().scope(): tpu_embedding_v2.TPUEmbedding( (tpu_embedding_v2_utils.FeatureConfig( table=tpu_embedding_v2_utils.TableConfig( name='table', vocabulary_size=4, dim=2, initializer=self.initializer,), name='watched'), tpu_embedding_v2_utils.FeatureConfig( table=tpu_embedding_v2_utils.TableConfig( name='table', vocabulary_size=4, dim=2, initializer=self.initializer), name='favorited')), tpu_embedding_v2_utils.SGD(learning_rate=0.1))
def test_feature_config_repr(self): table = tpu_embedding_v2_utils.TableConfig( vocabulary_size=2, dim=4, initializer=None, combiner='sum', name='table') feature_config = tpu_embedding_v2_utils.FeatureConfig( table=table, name='feature') self.assertEqual( repr(feature_config), 'FeatureConfig(table=TableConfig(vocabulary_size=2, dim=4, ' 'initializer=None, optimizer=None, combiner=\'sum\', name=\'table\'), ' 'max_sequence_length=0, validate_weights_and_indices=True, ' 'name=\'feature\')')
def tpu_embedding_config(): feature_configs = [] for dim, vocab, name in table_data: feature_configs.append(tpu_embedding_v2_utils.FeatureConfig( table=tpu_embedding_v2_utils.TableConfig( vocabulary_size=int(vocab), dim=int(dim), initializer=init_ops_v2.Zeros(), name=name))) optimizer = tpu_embedding_v2_utils.Adagrad( learning_rate=0.1) with strategy.scope(): mid_level_api = tpu_embedding_v2.TPUEmbedding( feature_config=feature_configs, optimizer=optimizer) mid_level_api._output_shapes = [TensorShape(128)] * len(feature_configs) return mid_level_api._create_config_proto()
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 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_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])
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) args = parser.parse_args() # embedding_values = np.array(list(range(32)), 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=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 feature_watched_values = np.random.randint(0, features, (batch * nnz * 2, )) # print("Feature: ", feature_watched_values) batch_size = batch * nnz # feature_watched_values = [0, 0, 1, 0, 1, 1] # feature_watched_row_lengths = [1, 2, 2, 1]
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 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 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.' )