def test_scope_reuse_sparse_embedding_lookup(self): indices = [ [0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [1, 0, 0], [1, 1, 0], [1, 1, 1], ] ids = [0, 1, -1, -1, 2, 0, 1] shape = [2, 3, 4] sparse_ids = sparse_tensor.SparseTensor( constant_op.constant(indices, dtypes.int64), constant_op.constant(ids, dtypes.int64), constant_op.constant(shape, dtypes.int64), ) with variable_scope.variable_scope("test", reuse=variable_scope.AUTO_REUSE): p1 = de.get_variable(name="p1") with variable_scope.variable_scope("q"): _, t1 = de.embedding_lookup_sparse(p1, sparse_ids, None, name="sp_emb", return_trainable=True) with variable_scope.variable_scope("test", reuse=variable_scope.AUTO_REUSE): p1_reuse = de.get_variable(name="p1") p2 = de.get_variable(name="p2") with variable_scope.variable_scope("q"): _, t2 = de.embedding_lookup_sparse(p2, sparse_ids, None, name="sp_emb", return_trainable=True) self.assertAllEqual(p1.name, "test/p1") self.assertAllEqual(p2.name, "test/p2") self.assertAllEqual(p1, p1_reuse) self.assertEqual(t1.name, "test/q/sp_emb/embedding_lookup/TrainableWrapper:0") self.assertEqual( t2.name, "test/q/sp_emb/embedding_lookup/TrainableWrapper_1:0") self.assertAllEqual(p1._tables[0].name, "test_p1_mht_1of1") self.assertAllEqual(p1_reuse._tables[0].name, "test_p1_mht_1of1") self.assertAllEqual(p2._tables[0].name, "test_p2_mht_1of1")
def loss_fn(emb, trainables): test_var, trainable = de.embedding_lookup_sparse( emb, sp_ids, sp_weights=None, combiner="sum", return_trainable=True) pred = math_ops.matmul(test_var, x) trainables.clear() trainables.append(trainable) return pred * pred
def test_embedding_lookup_sparse_shape_checking(self): with self.session(use_gpu=test_util.is_gpu_available(), config=default_config): embed_dim = 4 embedding_weights_nn = variable_scope.get_variable( "n", shape=[100, embed_dim], use_resource=False) embedding_weights_de = _random_weights(embed_dim=4) sparse_ids, _ = ids_and_weights_3d(embed_dim=embed_dim) embedding_lookup_base = embedding_ops.embedding_lookup_sparse( embedding_weights_nn, sparse_ids, None) embedding_lookup_test = de.embedding_lookup_sparse( embedding_weights_de, sparse_ids, None) self.assertTrue(embedding_lookup_base.get_shape().as_list() == embedding_lookup_test.get_shape().as_list())
def test_embedding_lookup_sparse(self): var_id = 0 for ( num_shards, initial_mode, combiner, k_dtype, d_dtype, ignore_weights, dim, ) in itertools.product( [1, 3], ["constant", "random"], [ "sum", "mean", "sqrtn", ], [dtypes.int64], [ dtypes.float32, ], [True, False], [1, 5], ): vocab_size = 2**31 if k_dtype == dtypes.int32 else 2**63 batch_size = 5 ( sp_ids, sp_weights, ids, weights, vals_per_batch_entry, ) = self._random_ids_and_weights(batch_size, vocab_size, k_dtype, d_dtype) grouped_ids = self._group_by_batch_entry(ids, vals_per_batch_entry) grouped_weights = self._group_by_batch_entry(weights, vals_per_batch_entry) grouped_ignored_weights = self._group_by_batch_entry( np.ones(np.sum(vals_per_batch_entry)), vals_per_batch_entry) with self.session(use_gpu=test_util.is_gpu_available(), config=default_config): var_id += 1 params = de.get_variable( "t1000-" + str(var_id), key_dtype=k_dtype, value_dtype=d_dtype, devices=_get_devices() * num_shards, initializer=1.0, dim=dim, ) random_init = params.lookup(ids) init_op = params.upsert(ids, random_init) self.evaluate(init_op) np_params = random_init.eval() grouped_params = self._group_by_batch_entry(np_params, vals_per_batch_entry) embedding_sum = de.embedding_lookup_sparse( params, sp_ids, None if ignore_weights else sp_weights, combiner=combiner, ) self.assertEqual(embedding_sum.dtype, d_dtype) tf_embedding_sum = embedding_sum.eval() np_embedding_sum, np_weight_sum, np_weight_sq_sum = embedding_result( grouped_params, grouped_ids, weight_vals=grouped_ignored_weights if ignore_weights else grouped_weights, ) if combiner == "mean": np_embedding_sum /= np.reshape(np_weight_sum, (batch_size, 1)) if combiner == "sqrtn": np_embedding_sum /= np.reshape(np.sqrt(np_weight_sq_sum), (batch_size, 1)) rtol = 1e-6 atol = rtol self.assertAllClose(np_embedding_sum, tf_embedding_sum, rtol, atol)
def common_minimize_trainable(self, base_opt, test_opt, name): base_opt = de.DynamicEmbeddingOptimizer(base_opt) test_opt = de.DynamicEmbeddingOptimizer(test_opt) id = 0 config = config_pb2.ConfigProto() config.allow_soft_placement = False for ( num_shards, k_dtype, d_dtype, initial_mode, dim, run_step, ) in itertools.product( [1, 2], [dtypes.int64], [ dtypes.float32, ], [ "constant", ], [1, 10], [10], ): id += 1 raw_init_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] raw_init_vals = [ [ x, ] * dim for x in [0.0, 0.1, 0.3, 0.8, 0.16, 0.25, 0.36, 0.49, 0.64, 0.81] ] raw_ids = constant_op.constant([1, 3, 3, 9], dtype=k_dtype) sp_ids = sparse_tensor.SparseTensor( indices=[ [0, 0], [0, 1], [1, 0], [2, 1], ], values=raw_ids, dense_shape=[3, 2], ) x = constant_op.constant([[_x * dim] for _x in [[0.4], [0.5], [0.6]]], dtype=d_dtype) x = array_ops.reshape(x, shape=(3 * dim, 1)) # base branch with self.session(use_gpu=test_util.is_gpu_available(), config=default_config) as sess: base_var = variables.Variable( np.array(raw_init_vals).reshape([len(raw_init_ids), dim]), dtype=d_dtype, shape=[len(raw_init_ids), dim], ) base_embedding = embedding_ops.embedding_lookup_sparse( base_var, sp_ids, None, combiner="sum") base_embedding = array_ops.reshape(base_embedding, shape=[1, 3 * dim]) pred0 = math_ops.matmul(base_embedding, x) loss0 = pred0 * pred0 base_opt_op = base_opt.minimize(loss0, var_list=[base_var]) # run base self.evaluate(variables.global_variables_initializer()) for _ in range(run_step): sess.run(base_opt_op) base_var_val = self.evaluate(base_var) # test branch with self.session(config=default_config, use_gpu=test_util.is_gpu_available()) as sess: # test var prepare embeddings = de.get_variable( "t1030-" + name + str(id), key_dtype=k_dtype, value_dtype=d_dtype, devices=_get_devices() * num_shards, initializer=1.0, dim=dim, ) self.device_check(embeddings) init_ids = constant_op.constant(raw_init_ids, dtype=k_dtype) init_vals = constant_op.constant(raw_init_vals, dtype=d_dtype) init_op = embeddings.upsert(init_ids, init_vals) self.evaluate(init_op) test_var, trainable = de.embedding_lookup_sparse( embeddings, sp_ids, sp_weights=None, combiner="sum", return_trainable=True, ) pred1 = math_ops.matmul( array_ops.reshape(test_var, shape=[1, 3 * dim]), x) loss1 = pred1 * pred1 test_opt_op = test_opt.minimize(loss1, var_list=[trainable]) self.evaluate(variables.global_variables_initializer()) self.assertAllCloseAccordingToType( np.array(raw_init_vals).reshape([len(raw_init_ids), dim]), self.evaluate(base_var), ) # Run `run_step` step of sgd for _ in range(run_step): sess.run(test_opt_op) if test_util.is_gpu_available(): self.assertTrue( _check_device(embeddings.tables[0].resource_handle, "GPU")) table_var_val = self.evaluate( array_ops.reshape(embeddings.lookup(init_ids), shape=[10, dim])) # Validate updated params self.assertAllCloseAccordingToType( base_var_val, table_var_val, msg="Cond:{},{},{},{},{}".format(num_shards, k_dtype, d_dtype, dim, run_step), )
def common_minimize_trainable(self, base_opt, test_opt, name): base_opt = de.DynamicEmbeddingOptimizer(base_opt) test_opt = de.DynamicEmbeddingOptimizer(test_opt) id = 0 for ( num_shards, k_dtype, d_dtype, initial_mode, dim, run_step, ) in itertools.product( [3], [dtypes.int64], [ dtypes.float32, ], [ "constant", ], [1, 10], [10], ): with ops.Graph().as_default(): id += 1 raw_init_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] raw_init_vals = [ [ x, ] * dim for x in [0.0, 0.1, 0.3, 0.8, 0.16, 0.25, 0.36, 0.49, 0.64, 0.81] ] raw_ids = constant_op.constant([1, 3, 3, 9], dtype=k_dtype) sp_ids = sparse_tensor.SparseTensor( indices=[ [0, 0], [0, 1], [1, 0], [2, 1], ], values=raw_ids, dense_shape=[3, 2], ) x = constant_op.constant([[_x * dim] for _x in [[0.4], [0.5], [0.6]]], dtype=d_dtype) x = array_ops.reshape(x, shape=(3 * dim, 1)) # base var prepare base_var = variables.Variable( np.array(raw_init_vals).reshape([len(raw_init_ids), dim]), dtype=d_dtype, shape=[len(raw_init_ids), dim], ) # test var prepare embeddings = de.get_variable( "t1030-" + name + str(id), key_dtype=k_dtype, value_dtype=d_dtype, devices=_get_devices() * num_shards, initializer=1.0, dim=dim, ) init_ids = constant_op.constant(raw_init_ids, dtype=k_dtype) init_vals = constant_op.constant(raw_init_vals, dtype=d_dtype) init_op = embeddings.upsert(init_ids, init_vals) # base branch base_embedding = embedding_ops.embedding_lookup_sparse( base_var, sp_ids, None, combiner="sum") base_embedding = array_ops.reshape(base_embedding, shape=[1, 3 * dim]) pred0 = math_ops.matmul(base_embedding, x) loss0 = pred0 * pred0 base_opt_op = base_opt.minimize(loss0, var_list=[base_var]) # test branch test_var, trainable = de.embedding_lookup_sparse( embeddings, sp_ids, sp_weights=None, combiner="sum", return_trainable=True, ) pred1 = math_ops.matmul( array_ops.reshape(test_var, shape=[1, 3 * dim]), x) loss1 = pred1 * pred1 gstep = training_util.create_global_step() test_opt_op = test_opt.minimize(loss1, var_list=[trainable], global_step=gstep) table_var = array_ops.reshape(embeddings.lookup(init_ids), shape=[10, dim]) with monitored_session.MonitoredTrainingSession( is_chief=True, config=default_config) as sess: sess.run(init_op) self.assertAllCloseAccordingToType( np.array(raw_init_vals).reshape( [len(raw_init_ids), dim]), sess.run(base_var), ) # run base for _ in range(run_step): sess.run(base_opt_op) sess.run(test_opt_op) # Validate global_step self.assertEqual(run_step, sess.run(gstep)) # Validate updated params self.assertAllCloseAccordingToType( sess.run(base_var), sess.run(table_var), msg="Cond:{},{},{},{},{}".format( num_shards, k_dtype, d_dtype, dim, run_step), ) self.device_check(embeddings)