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)