def create_status(self, **kwargs):
"""
Create relative frequency status variables.
"""
for tw in self.var.trainable_wrappers:
slots = [
self.opt.get_slot(tw, name)
for name in self.opt.get_slot_names()
]
self.var_in_slots += [x.params for x in slots if x is not None]
scope = variable_scope.get_variable_scope()
if scope.name:
scope_name = scope.name + '/frequency_status'
else:
scope_name = 'frequency_status'
status_name = self.var.name + '/' + self.opt.get_name()
with ops.name_scope(scope_name, "frequency_status",
[]) as unique_scope:
full_name = unique_scope + '/' + status_name
self.freq_var = dynamic_embedding_ops.get_variable(
key_dtype=self.var.key_dtype,
value_dtype=dtypes.int32,
dim=1,
name=full_name,
devices=self.var.devices,
partitioner=self.var.partition_fn,
initializer=self.default_count,
trainable=False,
)
def test_fn():
embeddings = deo.get_variable('t2020-v2-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
dim=dim)
self.device_check(embeddings)
trainables = []
init_ids = constant_op.constant(raw_init_ids, dtype=k_dtype)
init_vals = constant_op.constant(raw_init_vals, dtype=d_dtype)
self.evaluate(embeddings.upsert(init_ids, init_vals))
def var_fn():
return trainables
def loss_fn(x, trainables):
ids = constant_op.constant(raw_ids, dtype=k_dtype)
pred, trainable = deo.embedding_lookup(
[x], ids, return_trainable=True)
del trainables[:]
trainables.append(trainable)
return pred * pred
test_opt_op = test_opt.minimize(
lambda: loss_fn(embeddings, trainables), var_fn)
self.evaluate(variables.global_variables_initializer())
for _ in range(run_step):
self.evaluate(test_opt_op)
return embeddings.lookup(init_ids)
def test_check_ops_number(self):
self.assertTrue(deo.get_model_mode() == "train")
deo.enable_inference_mode()
self.assertTrue(deo.get_model_mode() == "inference")
deo.enable_train_mode()
self.assertTrue(deo.get_model_mode() == "train")
for fn, assign_num, read_num in [(deo.enable_train_mode, 1, 2),
(deo.enable_inference_mode, 0, 1)]:
fn()
embeddings = deo.get_variable('ModeModeTest' + str(assign_num),
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
devices=_get_devices(),
initializer=1.,
dim=8)
ids = constant_op.constant([0, 1, 2, 3, 4], dtype=dtypes.int64)
test_var, trainable = deo.embedding_lookup([embeddings],
ids,
return_trainable=True)
_ = math_ops.add(test_var, 1)
op_list = ops.get_default_graph().get_operations()
op_list_assign = [
op.name for op in op_list
if "AssignBeforeReadVariable" in op.name
]
op_list_read = [
op.name for op in op_list if "ReadVariableOp" in op.name
]
self.assertTrue(len(op_list_assign) == assign_num)
self.assertTrue(len(op_list_read) == read_num)
ops.reset_default_graph()
def test_check_find_ops_number(self):
self.assertTrue(deo.get_update_mode() == "convergence_priority")
deo.enable_speed_priority()
self.assertTrue(deo.get_update_mode() == "speed_priority")
deo.enable_convergence_priority()
self.assertTrue(deo.get_update_mode() == "convergence_priority")
for fn, nm in [(deo.enable_speed_priority, 2),
(deo.enable_convergence_priority, 6)]:
fn()
embeddings = deo.get_variable('UpdateModeTest' + str(nm),
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
devices=_get_devices(),
initializer=1.,
dim=8)
ids = constant_op.constant([0, 1, 2, 3, 4], dtype=dtypes.int64)
test_var, trainable = deo.embedding_lookup([embeddings],
ids,
return_trainable=True)
pred = math_ops.add(test_var, 1)
loss = pred * pred
opt = adagrad.AdagradOptimizer(0.1)
opt.minimize(loss, var_list=[trainable])
op_list = ops.get_default_graph().get_operations()
op_list = [op.name for op in op_list if "Find" in op.name]
self.assertTrue(len(op_list) == nm)
ops.reset_default_graph()
def test_ops_with_var_and_ftrl(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
ftrl.FtrlOptimizer(0.1),
]
self.common_run_context(var_list, opt_list, name='ftrl_test')
def test_ops_with_var_and_lamb(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
lamb.LAMBOptimizer(),
]
self.common_run_context(var_list, opt_list, name='lamb_test')
def test_ops_with_var_and_momentum(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
momentum.MomentumOptimizer(0.1, 0.1),
]
self.common_run_context(var_list, opt_list, name='momentum_test')
def test_ops_with_var_and_pgd(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
pgd.ProximalGradientDescentOptimizer(0.1),
]
self.common_run_context(var_list, opt_list, name='pgd_test')
def test_ops_with_var_and_rmsprop(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
rmsprop.RMSPropOptimizer(0.1),
]
self.common_run_context(var_list, opt_list, name='rmsprop_test')
def test_ops_with_var_and_adam(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
adam.AdamOptimizer(0.1),
]
self.common_run_context(var_list, opt_list, name='adam_test')
def test_ops_with_var_and_adagrad_da(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
gstep = training_util.create_global_step()
opt_list = [
adagrad_da.AdagradDAOptimizer(0.1, gstep),
]
self.common_run_context(var_list, opt_list, name='adagrad_da_test')
def _add_and_initialize_devar(self, name, keys, values, dim):
var = deo.get_variable(name,
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
dim=dim)
self.evaluate(
var.upsert(constant_op.constant(keys, dtypes.int64),
constant_op.constant(values, dtypes.float32)))
return var
def test_init_exception_invalid_policy(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt = gradient_descent.GradientDescentOptimizer(0.1)
err = None
with self.assertRaises(TypeError):
_ = dvr.VariableRestrictor(var_list=var_list,
optimizer_list=[opt],
policy=None)
def test_ops_with_var_and_gradient_descent(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
gradient_descent.GradientDescentOptimizer(0.1),
]
self.common_run_context(var_list,
opt_list,
name='gradient_descent_test')
def test_ops_with_vars_and_optimizers(self):
num_vars = 4
var_list = [
deo.get_variable('sp_var_' + str(_i), initializer=0.0, dim=2)
for _i in range(num_vars)
]
opt_list = get_multiple_optimizers()
self.common_run_context(var_list,
opt_list,
name='multiple_optimizer_test')
def test_ops_with_var_and_proximal_adagrad(self):
var_list = [
deo.get_variable('sp_var', initializer=0.0, dim=2),
]
opt_list = [
proximal_adagrad.ProximalAdagradOptimizer(0.1),
]
self.common_run_context(var_list,
opt_list,
name='proximal_adagrad_test')
def test_ops_with_various_variables_and_optimizers(self):
key_dtypes = [dtypes.int64]
n = 0
var_list = list()
for _kt in key_dtypes:
_var = deo.get_variable('sp_var_' + str(n),
initializer=0.0,
key_dtype=_kt,
dim=8)
var_list.append(_var)
n += 1
opt_list = get_multiple_optimizers()
self.common_run_context(var_list, opt_list, name='cross_context_test')
def __init__(self,
key_dtype,
value_dtype,
dim,
num_shards,
var_name,
devar_name,
run_id,
x=None):
self.key_dtype = key_dtype
self.value_dtype = value_dtype
self.dim = dim
# common define
init_ids = [0, 1, 2]
init_vals = np.random.rand(3, self.dim)
raw_ids = [1]
if x is None:
self.x = constant_op.constant(np.random.rand(
self.dim, len(raw_ids)),
dtype=self.value_dtype)
else:
self.x = ops.convert_to_tensor(x, dtype=self.value_dtype)
# variable graph
self.var = resource_variable_ops.ResourceVariable(
name='t2020-' + var_name + str(run_id),
initial_value=init_vals,
dtype=self.value_dtype)
ids = constant_op.constant(raw_ids, dtype=self.key_dtype)
self.var_lookup = embedding_ops.embedding_lookup([self.var], ids)
self.var_pred = math_ops.matmul(self.var_lookup, self.x)
self.var_loss = self.var_pred * self.var_pred
self.var_opt_op = adam.AdamOptimizer(1.0).minimize(self.var_loss)
# deo variable graph
self.devar = deo.get_variable(name='t2020-' + devar_name + str(run_id),
key_dtype=self.key_dtype,
value_dtype=self.value_dtype,
devices=_get_devices() * num_shards,
initializer=1.,
dim=dim)
self.devar_init_op = self.devar.upsert(
constant_op.constant(init_ids, dtype=self.key_dtype),
constant_op.constant(init_vals, dtype=self.value_dtype))
self.devar_lookup, _ = deo.embedding_lookup([self.devar],
ids,
return_trainable=True)
self.devar_pred = math_ops.matmul(self.devar_lookup, self.x)
self.devar_loss = self.devar_pred * self.devar_pred
self.devar_opt_op = adam.AdamOptimizer(1.0).minimize(self.devar_loss)
def test_inference_numberic_correctness(self):
train_pred = None
infer_pred = None
dim = 8
initializer = init_ops.random_normal_initializer(0.0, 0.001)
raw_init_vals = np.random.rand(100, dim)
for fn in [deo.enable_train_mode, deo.enable_inference_mode]:
with ops.Graph().as_default():
fn()
init_ids = constant_op.constant(list(range(100)),
dtype=dtypes.int64)
init_vals = constant_op.constant(raw_init_vals,
dtype=dtypes.float32)
with variable_scope.variable_scope(
"modelmode", reuse=variable_scope.AUTO_REUSE):
embeddings = deo.get_variable('ModelModeTest-numberic',
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
devices=_get_devices() * 2,
initializer=initializer,
dim=dim)
w = variables.Variable(1.0, name="w")
_ = training_util.create_global_step()
init_op = embeddings.upsert(init_ids, init_vals)
ids = constant_op.constant([0, 1, 2, 3, 4], dtype=dtypes.int64)
test_var, trainable = deo.embedding_lookup(
[embeddings], ids, return_trainable=True)
pred = math_ops.add(test_var, 1) * w
loss = pred * pred
opt = adagrad.AdagradOptimizer(0.1)
opt.minimize(loss)
with monitored_session.MonitoredTrainingSession(
is_chief=True, config=default_config) as sess:
if deo.get_model_mode() == deo.ModelMode.TRAIN:
sess.run(init_op)
train_pred = sess.run(pred)
elif deo.get_model_mode() == deo.ModelMode.INFERENCE:
sess.run(init_op)
infer_pred = sess.run(pred)
ops.reset_default_graph()
self.assertAllEqual(train_pred, infer_pred)
def create_graph_and_policy(optimizer, ids, mode):
var = deo.get_variable('sp_var',
key_dtype=ids.dtype,
value_dtype=dtypes.float32,
initializer=-1.,
dim=2)
embedding_w, trainable_wrapper = deo.embedding_lookup(
var, ids, return_trainable=True)
loss = math_ops.reduce_sum(embedding_w)
train_op = optimizer.minimize(loss, var_list=[trainable_wrapper])
if mode == 'timestamp':
policy = dvr.TimestampRestrictPolicy(var, optimizer)
elif mode == 'frequency':
policy = dvr.FrequencyRestrictPolicy(var, optimizer)
else:
raise NotImplementedError
slot_vars = select_slot_vars(trainable_wrapper, optimizer)
return slot_vars, policy, var, train_op
def create_status(self, **kwargs):
"""
Create relative probability status variables.
"""
scope = variable_scope.get_variable_scope()
if scope.name:
scope_name = scope.name + '/probability_status_for_filter'
else:
scope_name = 'probability_status_for_filter'
with ops.name_scope(scope_name, "probability_status_for_filter",
[]) as unique_scope:
full_name = unique_scope + '/' + self.var.name
self.tstp_var = dynamic_embedding_ops.get_variable(
key_dtype=self.var.key_dtype,
value_dtype=dtypes.int32,
dim=1,
name=full_name,
devices=self.var.devices,
partitioner=self.var.partition_fn,
initializer=self.default_tstp,
trainable=False,
)
def test_fn():
embeddings = deo.get_variable('s6030-v2-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
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)
self.evaluate(embeddings.upsert(init_ids, init_vals))
trainables = []
def var_fn():
return trainables
def loss_fn(emb, trainables):
test_var, trainable = deo.safe_embedding_lookup_sparse(
emb,
sp_ids_de,
sparse_weights=None,
combiner="sum",
return_trainable=True)
pred = math_ops.matmul(test_var, x)
del trainables[:]
trainables.append(trainable)
return pred * pred
test_opt_op = test_opt.minimize(
lambda: loss_fn(embeddings, trainables), var_fn)
self.evaluate(variables.global_variables_initializer())
for _ in range(run_step):
self.evaluate(test_opt_op)
return embeddings.lookup(init_ids)
def _add_devar(self, name, dim):
var = deo.get_variable(name,
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
dim=dim)
return var
def test_training_save_restore(self):
# embedding_lookup does not work in eager mode when num_shards is more than 1.
ops.disable_eager_execution()
keys_type_list = [dtypes.int64] if test_util.is_gpu_available() else [
dtypes.int64, dtypes.string
]
for run_id, num_shards, key_dtype, value_dtype, _, dim, run_step \
in _next_run_step_config(keys_type_list):
save_dir = os.path.join(self.get_temp_dir(), "save_restore")
save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
ids = script_ops.py_func(
_create_dynamic_shape_tensor(dtype=_type_converter(key_dtype)),
inp=[],
Tout=key_dtype,
stateful=True)
params = deo.get_variable(
name="params-test-0915-" + str(run_id),
key_dtype=key_dtype,
value_dtype=value_dtype,
devices=_get_devices() * num_shards,
initializer=init_ops.random_normal_initializer(0.0, 0.01),
dim=dim)
_, var0 = deo.embedding_lookup(params, ids, return_trainable=True)
loss = lambda: var0 * var0
params_keys, params_vals = params.export()
opt = adam.AdamOptimizer(0.3)
mini = opt.minimize(loss, var_list=[var0])
opt_slots = [opt.get_slot(var0, _s) for _s in opt.get_slot_names()]
_saver = saver_lib.Saver([params] +
[_s.params for _s in opt_slots])
with self.session(config=default_config,
use_gpu=test_util.is_gpu_available()) as sess:
self.evaluate(variables.global_variables_initializer())
for _i in range(run_step):
self.evaluate([mini])
size_before_saved = self.evaluate(params.size())
np_params_keys_before_saved = self.evaluate(params_keys)
np_params_vals_before_saved = self.evaluate(params_vals)
opt_slots_kv_pairs = [_s.params.export() for _s in opt_slots]
np_slots_kv_pairs_before_saved = [
self.evaluate(_kv) for _kv in opt_slots_kv_pairs
]
_saver.save(sess, save_path)
with self.session(config=default_config,
use_gpu=test_util.is_gpu_available()) as sess:
self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(0, self.evaluate(params.size()))
_saver.restore(sess, save_path)
params_keys_restored, params_vals_restored = params.export()
size_after_restored = self.evaluate(params.size())
np_params_keys_after_restored = self.evaluate(
params_keys_restored)
np_params_vals_after_restored = self.evaluate(
params_vals_restored)
opt_slots_kv_pairs_restored = [
_s.params.export() for _s in opt_slots
]
np_slots_kv_pairs_after_restored = [
self.evaluate(_kv) for _kv in opt_slots_kv_pairs_restored
]
self.assertAllEqual(size_before_saved, size_after_restored)
self.assertAllEqual(np.sort(np_params_keys_before_saved),
np.sort(np_params_keys_after_restored))
self.assertAllEqual(
np.sort(np_params_vals_before_saved, axis=0),
np.sort(np_params_vals_after_restored, axis=0))
for pairs_before, pairs_after in zip(
np_slots_kv_pairs_before_saved,
np_slots_kv_pairs_after_restored):
self.assertAllEqual(np.sort(pairs_before[0], axis=0),
np.sort(pairs_after[0], axis=0))
self.assertAllEqual(np.sort(pairs_before[1], axis=0),
np.sort(pairs_after[1], axis=0))
if test_util.is_gpu_available():
self.assertTrue(
_check_device(params.tables[0].resource_handle, "GPU"))
def common_run_context(self, var_list, opt_list, name):
batch_size = 2
sample_length = 3
emb_domain_list = list()
tws = list()
cluster = ps_worker_cluster(ps_num=2)
ps_servers, worker_servers, cluster_def = cluster
config = config_pb2.ConfigProto(
cluster_def=cluster_def,
experimental=config_pb2.ConfigProto.Experimental(
share_session_state_in_clusterspec_propagation=True, ),
allow_soft_placement=False,
inter_op_parallelism_threads=2,
intra_op_parallelism_threads=2,
gpu_options=config_pb2.GPUOptions(allow_growth=True),
)
dev_placement = device_setter.replica_device_setter(
ps_tasks=2,
ps_device='/job:ps',
worker_device='/job:worker',
cluster=cluster_def,
)
with ops.device(dev_placement):
shared_var_0 = deo.get_variable('distributed_sp_var_0',
initializer=0.0,
devices=['/job:worker/task:0'],
dim=8)
shared_var_1 = deo.get_variable('distributed_sp_var_1',
initializer=0.0,
devices=['/job:worker/task:0'],
dim=4)
opt_list = get_multiple_optimizers()
distributed_var_list = [shared_var_0, shared_var_1]
for _v in distributed_var_list:
ids = random_ops.random_uniform((batch_size, sample_length),
maxval=1000000,
dtype=_v.key_dtype)
ids = array_ops.reshape(ids, (-1, ))
_, tw = deo.embedding_lookup(_v, ids, return_trainable=True)
tws.append(tw)
_collapse = array_ops.reshape(tw, (batch_size, -1))
_logits = math_ops.reduce_sum(_collapse, axis=1)
_logits = math_ops.cast(_logits, dtypes.float32)
emb_domain_list.append(_logits)
logits = math_ops.add_n(emb_domain_list)
labels = array_ops.zeros((batch_size, ), dtype=dtypes.float32)
loss = math_ops.reduce_mean(
nn_impl.sigmoid_cross_entropy_with_logits(
logits=logits,
labels=labels,
))
_train_ops = list()
for _opt in opt_list:
_train_ops.append(_opt.minimize(loss))
train_op = control_flow_ops.group(_train_ops)
restrictor = dvr.VariableRestrictor(var_list=distributed_var_list,
optimizer_list=opt_list)
update_op = restrictor.update()
threshold = int(batch_size * sample_length * 1.5)
factor = 1.2
restrict_op = restrictor.restrict(threshold=threshold,
factor=factor)
policies = list(itertools.chain(*restrictor.policy_group.values()))
tstp_vars = [policy.tstp_var for policy in policies]
slot_vars = list()
for tw in tws:
for opt in opt_list:
slot_vars += select_slot_vars(tw, opt)
with session.Session(worker_servers[0].target, config=config) as sess:
sess.run(variables.global_variables_initializer())
n, MAX_ITER = 0, 1000
while n < MAX_ITER:
sess.run([train_op, update_op])
if all(
sess.run(var.size()) > threshold * factor
for var in distributed_var_list):
break
s1 = sess.run([var.size() for var in distributed_var_list])
s2 = sess.run([tv.size() for tv in tstp_vars])
s3 = sess.run([sv.size() for sv in slot_vars])
self.assertAllGreater(s1, threshold * factor)
self.assertAllGreater(s2, threshold * factor)
if s3:
self.assertAllGreater(s3, threshold * factor)
sess.run(restrict_op)
s1 = sess.run([var.size() for var in distributed_var_list])
s2 = sess.run([tv.size() for tv in tstp_vars])
s3 = sess.run([sv.size() for sv in slot_vars])
self.assertAllLess(s1, threshold * factor + 1)
self.assertAllLess(s2, threshold * factor + 1)
if s3:
self.assertAllLess(s3, threshold * factor + 1)
sess.close()
def common_minimize_trainable(self, base_opt, test_opt, name):
deo.enable_train_mode()
for run_id, num_shards, k_dtype, d_dtype, initial_mode, dim, run_step \
in _next_run_step_config():
with self.session(use_gpu=test_util.is_gpu_available(),
config=default_config) as sess:
# common define
raw_init_ids = [0, 1]
raw_init_vals = np.random.rand(2, dim)
raw_ids = [
0,
]
x = constant_op.constant(np.random.rand(dim, len(raw_ids)),
dtype=d_dtype)
# base graph
base_var = resource_variable_ops.ResourceVariable(
raw_init_vals, dtype=d_dtype)
ids = constant_op.constant(raw_ids, dtype=k_dtype)
pred0 = math_ops.matmul(
embedding_ops.embedding_lookup([base_var], ids), x)
loss0 = pred0 * pred0
base_opt_op = base_opt.minimize(loss0)
# test graph
embeddings = deo.get_variable('t2020-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
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 = deo.embedding_lookup(
[embeddings], ids, return_trainable=True)
pred1 = math_ops.matmul(test_var, x)
loss1 = pred1 * pred1
test_opt_op = test_opt.minimize(loss1, var_list=[trainable])
self.evaluate(variables.global_variables_initializer())
for _ in range(run_step):
sess.run(base_opt_op)
# Fetch params to validate initial values
self.assertAllCloseAccordingToType(raw_init_vals[raw_ids],
self.evaluate(test_var))
# Run `run_step` step of sgd
for _ in range(run_step):
sess.run(test_opt_op)
table_var = embeddings.lookup(ids)
# Validate updated params
self.assertAllCloseAccordingToType(
self.evaluate(base_var)[raw_ids],
self.evaluate(table_var),
msg="Cond:{},{},{},{},{},{}".format(
num_shards, k_dtype, d_dtype, initial_mode, dim,
run_step))
def common_minimize_trainable(self, base_opt, test_opt, name):
config = config_pb2.ConfigProto()
deo.enable_train_mode()
config.allow_soft_placement = False
for run_id, num_shards, k_dtype, d_dtype, initial_mode, dim, run_step \
in _next_run_step_config():
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 = deo.get_variable('t1030-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
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 = deo.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):
if test_util.is_gpu_available():
keys_type_list = [dtypes.int64]
else:
keys_type_list = [dtypes.int64, dtypes.string]
deo.enable_train_mode()
config = config_pb2.ConfigProto(
allow_soft_placement=True,
gpu_options=config_pb2.GPUOptions(allow_growth=True))
for run_id, num_shards, k_dtype, d_dtype, initial_mode, dim, run_step \
in _next_run_step_config(keys_type_list):
with self.session(config=config,
use_gpu=test_util.is_gpu_available()):
raw_init_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
if k_dtype == dtypes.string:
raw_init_ids = [str(i) for i in raw_init_ids]
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_py = [1, 3, 3, 9]
raw_ids_nn = constant_op.constant(raw_ids_py,
dtype=dtypes.int64)
sp_ids_nn = sparse_tensor.SparseTensor(indices=[
[0, 0],
[0, 1],
[1, 0],
[2, 1],
],
values=raw_ids_nn,
dense_shape=[3, 2])
if k_dtype != dtypes.string:
raw_ids_de = raw_ids_nn
else:
raw_ids_de = constant_op.constant(
[str(i) for i in raw_ids_py], dtype=k_dtype)
sp_ids_de = sparse_tensor.SparseTensor(indices=[
[0, 0],
[0, 1],
[1, 0],
[2, 1],
],
values=raw_ids_de,
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])
base_embedding = embedding_ops.safe_embedding_lookup_sparse(
base_var, sp_ids_nn, 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 var prepare
embeddings = deo.get_variable('s6030-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
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 branch
test_var, trainable = deo.safe_embedding_lookup_sparse(
embeddings,
sp_ids_de,
sparse_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 base
for _ in range(run_step):
self.evaluate(base_opt_op)
# Run `run_step` step of sgd
for _ in range(run_step):
self.evaluate(test_opt_op)
table_var = array_ops.reshape(embeddings.lookup(init_ids),
shape=[10, dim])
# Validate updated params
self.assertAllCloseAccordingToType(
self.evaluate(base_var),
self.evaluate(table_var),
msg="Cond:{},{},{},{},{}".format(num_shards, k_dtype,
d_dtype, dim, run_step))
def common_minimize_trainable(self, base_opt, test_opt, name):
if test_util.is_gpu_available():
keys_type_list = [dtypes.int64]
else:
keys_type_list = [dtypes.int64, dtypes.string]
deo.enable_train_mode()
for run_id, num_shards, k_dtype, d_dtype, initial_mode, dim, run_step \
in _next_run_step_config(keys_type_list):
with ops.Graph().as_default():
raw_init_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
if k_dtype == dtypes.string:
raw_init_ids = [str(i) for i in raw_init_ids]
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_py = [1, 3, 3, 9]
raw_ids_nn = constant_op.constant(raw_ids_py,
dtype=dtypes.int64)
raw_ids_de = raw_ids_nn
if k_dtype == dtypes.string:
raw_ids_de = constant_op.constant(
[str(i) for i in raw_ids_py], dtype=k_dtype)
sp_ids_nn = sparse_tensor.SparseTensor(indices=[
[0, 0],
[0, 1],
[1, 0],
[2, 1],
],
values=raw_ids_nn,
dense_shape=[3, 2])
sp_ids_de = sparse_tensor.SparseTensor(indices=[
[0, 0],
[0, 1],
[1, 0],
[2, 1],
],
values=raw_ids_de,
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 = deo.get_variable('t1030-' + name + str(run_id),
key_dtype=k_dtype,
value_dtype=d_dtype,
devices=_get_devices() *
num_shards,
initializer=1.,
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_nn, 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 = deo.embedding_lookup_sparse(
embeddings,
sp_ids_de,
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)
def build_distributed_graph():
batch_size = 4
shape_0 = [batch_size, 5]
shape_1 = [batch_size, 6]
maxval = int(0x7FFF)
server0 = server_lib.Server.create_local_server()
server1 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server0.target[len('grpc://'):]
job.tasks[1] = server1.target[len('grpc://'):]
config = config_pb2.ConfigProto(
cluster_def=cluster_def,
experimental=config_pb2.ConfigProto.Experimental(
share_session_state_in_clusterspec_propagation=True, ),
)
config.allow_soft_placement = False
with ops.device('/job:worker/task:0'):
feat_0 = random_ops.random_uniform(shape_0,
maxval=maxval,
dtype=dtypes.int64)
feat_0 = array_ops.reshape(feat_0, (-1, ))
feat_1 = random_ops.random_uniform(shape_1,
maxval=maxval,
dtype=dtypes.int64)
feat_1 = array_ops.reshape(feat_1, (-1, ))
var_0 = deo.get_variable(
name='sp_var_0',
devices=[
'/job:worker/task:1',
],
initializer=init_ops.random_normal_initializer(0, 0.005),
)
var_1 = deo.get_variable(
name='sp_var_1',
devices=[
'/job:worker/task:1',
],
initializer=init_ops.random_normal_initializer(0, 0.005),
)
var_list = [var_0, var_1]
_, tw_0 = deo.embedding_lookup(
params=var_0,
ids=feat_0,
name='sp_emb_0',
return_trainable=True,
)
_, tw_1 = deo.embedding_lookup(
params=var_1,
ids=feat_1,
name='sp_emb_1',
return_trainable=True,
)
collapse_0 = array_ops.reshape(tw_0, (batch_size, -1))
collapse_1 = array_ops.reshape(tw_1, (batch_size, -1))
logits_0 = math_ops.reduce_sum(collapse_0, axis=1)
logits_1 = math_ops.reduce_sum(collapse_1, axis=1)
logits = math_ops.add(logits_0, logits_1)
labels = array_ops.zeros((batch_size, ), dtype=dtypes.float32)
loss = math_ops.reduce_mean(
nn_impl.sigmoid_cross_entropy_with_logits(
logits=logits,
labels=labels,
))
optimizers = get_multiple_optimizers()
return server0, server1, config, var_list, optimizers, loss
