def test_scope_reuse_safe_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.safe_embedding_lookup_sparse(p1,
sparse_ids,
None,
name="safe_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.safe_embedding_lookup_sparse(p2,
sparse_ids,
None,
name="safe_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/safe_sp_emb/embedding_lookup_sparse/embedding_lookup/TrainableWrapper:0",
)
self.assertEqual(
t2.name,
"test/q/safe_sp_emb/embedding_lookup_sparse/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 test_safe_embedding_lookup_sparse_partitioned(self):
with self.session(use_gpu=test_util.is_gpu_available(),
config=default_config):
dim = 4
embedding_weights = _random_weights(embed_dim=dim, num_shards=3)
sparse_ids, sparse_weights = ids_and_weights_2d(embed_dim=dim)
valid_ids = np.array([0, 1, 2, -1])
# init
embedding_weights_values = embedding_weights.lookup(valid_ids).eval()
self.evaluate(
embedding_weights.upsert(valid_ids, embedding_weights_values))
embedding_lookup_result = de.safe_embedding_lookup_sparse(
embedding_weights, sparse_ids, None).eval()
self.assertAllClose(
embedding_lookup_result,
[
(embedding_weights_values[0] + embedding_weights_values[1] +
embedding_weights_values[3]) / 3.0,
embedding_weights_values[3],
[0] * 4,
embedding_weights_values[2],
(embedding_weights_values[0] + embedding_weights_values[1]) / 2.0,
],
)
def test_safe_embedding_lookup_sparse_return_special_vector(self):
with self.session(use_gpu=test_util.is_gpu_available(),
config=default_config):
dim = 4
embedding_weights = _random_weights(embed_dim=dim)
sparse_ids, sparse_weights = ids_and_weights_2d(embed_dim=dim)
valid_ids = np.array([0, 1, 2, 3, -1])
# init
embedding_weights_values = embedding_weights.lookup(valid_ids).eval()
self.evaluate(
embedding_weights.upsert(valid_ids, embedding_weights_values))
# check
embedding_lookup_result = de.safe_embedding_lookup_sparse(
embedding_weights, sparse_ids, sparse_weights, default_id=3).eval()
self.assertAllClose(
embedding_lookup_result,
[
(1.0 * embedding_weights_values[0] +
2.0 * embedding_weights_values[1] +
1.0 * embedding_weights_values[4]) / 4.0,
embedding_weights_values[4],
embedding_weights_values[3],
embedding_weights_values[2],
embedding_weights_values[3],
],
)
def test_safe_embedding_lookup_sparse_3d_no_weights(self):
with self.session(use_gpu=test_util.is_gpu_available(),
config=default_config):
embedding_weights = _random_weights()
sparse_ids, _ = ids_and_weights_3d()
valid_ids = np.array([0, 1, 2, -1])
# init
embedding_weights_values = embedding_weights.lookup(valid_ids).eval()
self.evaluate(
embedding_weights.upsert(valid_ids, embedding_weights_values))
embedding_lookup_result = de.safe_embedding_lookup_sparse(
embedding_weights, sparse_ids, None).eval()
self.assertAllClose(
embedding_lookup_result,
[
[
(embedding_weights_values[0] + embedding_weights_values[1] +
embedding_weights_values[3]) / 3.0,
embedding_weights_values[3],
[0] * 4,
],
[
embedding_weights_values[2],
(embedding_weights_values[0] + embedding_weights_values[1]) /
2.0,
[0] * 4,
],
],
)
def loss_fn(emb, trainables):
test_var, trainable = de.safe_embedding_lookup_sparse(
emb,
sp_ids,
sparse_weights=None,
combiner="sum",
return_trainable=True)
pred = math_ops.matmul(test_var, x)
trainables.clear()
trainables.append(trainable)
return pred * pred
def test_safe_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.safe_embedding_lookup_sparse(
embedding_weights_nn, sparse_ids, None)
embedding_lookup_test = de.safe_embedding_lookup_sparse(
embedding_weights_de, sparse_ids, None)
self.assertAllEqual(embedding_lookup_base.shape,
embedding_lookup_test.shape)
self.assertAllEqual(embedding_lookup_base.get_shape(),
embedding_lookup_test.get_shape())
def model_fn(features, labels, mode, params):
#logging.info('mode: %s, labels: %s, params: %s, features: %s', mode, labels, params, features)
if params["args"].get("addon_embedding"):
import tensorflow_recommenders_addons as tfra
import tensorflow_recommenders_addons.dynamic_embedding as dynamic_embedding
else:
import tensorflow.dynamic_embedding as dynamic_embedding
features.update(labels)
logging.info("------ build hyper parameters -------")
embedding_size = params["parameters"]["embedding_size"]
learning_rate = params["parameters"]["learning_rate"]
use_bn = params["parameters"]["use_bn"]
feat = params['features']
sparse_feat_list = list(set(feat["sparse"]) -
set(SPARSE_MASK)) if 'sparse' in feat else []
sparse_seq_feat_list = list(
set(feat["sparse_seq"]) -
set(SPARSE_SEQ_MASK)) if 'sparse_seq' in feat else []
sparse_seq_feat_list = []
#dense_feat_list = list(set(feat["dense"]) - set(DENSE_MASK)) if 'dense' in feat else []
# hashtable v1/v2 image均无法同时关bn和mask dense_feat
dense_feat_list = []
dense_seq_feat_list = list(
set(feat["dense_seq"]) -
set(DENSE_SEQ_MASK)) if 'dense_seq' in feat else []
sparse_feat_num = len(sparse_feat_list)
sparse_seq_num = len(sparse_seq_feat_list)
dense_feat_num = len(dense_feat_list)
dense_seq_feat_num = len(dense_seq_feat_list)
all_features = (sparse_feat_list + sparse_seq_feat_list + dense_feat_list +
dense_seq_feat_list)
batch_size = tf.shape(features[goods_id_feat])[0]
logging.info("------ show batch_size: {} -------".format(batch_size))
level_0_feats = list(
set(params.get('level_0_feat_list')) & set(all_features))
logging.info('level_0_feats: {}'.format(level_0_feats))
new_features = dict()
if params["args"].get("level_flag") and params["args"].get(
"job_type") == "export":
for feature_name in features:
if feature_name in level_0_feats:
new_features[feature_name] = tf.reshape(
tf.tile(tf.reshape(features[feature_name], [1, -1]),
[batch_size, 1]), [batch_size, -1])
else:
new_features[feature_name] = features[feature_name]
features = new_features
l2_reg = params["parameters"]["l2_reg"]
is_training = True if mode == tf.estimator.ModeKeys.TRAIN else False
has_label = True if 'is_imp' in features else False
logging.info("is_training: {}, has_label: {}, features: {}".format(
is_training, has_label, features))
logging.info("------ build embedding -------")
# def partition_fn(keys, shard_num=params["parameters"]["ps_nums"]):
# return tf.cast(keys % shard_num, dtype=tf.int32)
if is_training:
devices_info = [
"/job:ps/replica:0/task:{}/CPU:0".format(i)
for i in range(params["parameters"]["ps_num"])
]
initializer = tf.compat.v1.truncated_normal_initializer(0.0, 1e-2)
else:
devices_info = [
"/job:localhost/replica:0/task:{}/CPU:0".format(0)
for i in range(params["parameters"]["ps_num"])
]
initializer = tf.compat.v1.zeros_initializer()
logging.info("------ dynamic_embedding devices_info is {}-------".format(
devices_info))
if mode == tf.estimator.ModeKeys.PREDICT:
dynamic_embedding.enable_inference_mode()
deep_dynamic_variables = dynamic_embedding.get_variable(
name="deep_dynamic_embeddings",
devices=devices_info,
initializer=initializer,
# partitioner=partition_fn,
dim=embedding_size,
trainable=is_training,
#init_size=INIT_SIZE
)
sparse_feat = None
sparse_unique_ids = None
if sparse_feat_num > 0:
logging.info("------ build sparse feature -------")
id_list = sorted(sparse_feat_list)
ft_sparse_idx = tf.concat(
[tf.reshape(features[str(i)], [-1, 1]) for i in id_list], axis=1)
sparse_unique_ids, sparse_unique_idx = tf.unique(
tf.reshape(ft_sparse_idx, [-1]))
sparse_weights = dynamic_embedding.embedding_lookup(
params=deep_dynamic_variables,
ids=sparse_unique_ids,
name="deep_sparse_weights")
if params["args"].get("zero_padding"):
sparse_weights = tf.reshape(sparse_weights, [-1, embedding_size])
sparse_weights = tf.where(
tf.not_equal(
tf.expand_dims(sparse_unique_ids, axis=1),
tf.zeros_like(tf.expand_dims(sparse_unique_ids, axis=1))),
sparse_weights, tf.zeros_like(sparse_weights))
sparse_weights = tf.gather(sparse_weights, sparse_unique_idx)
sparse_feat = tf.reshape(
sparse_weights,
shape=[batch_size, sparse_feat_num * embedding_size])
sparse_seq_feat = None
sparse_seq_unique_ids = None
if sparse_seq_num > 0:
logging.info("---- build sparse seq feature ---")
if params["args"].get("merge_sparse_seq"):
sparse_seq_name_list = sorted(
sparse_seq_feat_list) #[B, s1], [B, s2], ... [B, sn]
ft_sparse_seq_ids = tf.concat(
[
tf.reshape(features[str(i)], [batch_size, -1])
for i in sparse_seq_name_list
],
axis=1) #[B, [s1, s2, ...sn]] => [B, per_seq_len*seq_num]
sparse_seq_unique_ids, sparse_seq_unique_idx = tf.unique(
tf.reshape(ft_sparse_seq_ids,
[-1])) #[u], [B*per_seq_len*seq_num]
sparse_seq_weights = dynamic_embedding.embedding_lookup(
params=deep_dynamic_variables,
ids=sparse_seq_unique_ids,
name="deep_sparse_seq_weights") #[u, e]
deep_embed_seq = tf.where(
tf.not_equal(
tf.expand_dims(sparse_seq_unique_ids, axis=1),
tf.zeros_like(tf.expand_dims(sparse_seq_unique_ids,
axis=1))), sparse_seq_weights,
tf.zeros_like(sparse_seq_weights)) #[u, e]
deep_embedding_seq = tf.reshape(
tf.gather(deep_embed_seq,
sparse_seq_unique_idx), #[B*per_seq_len*seq_num, e]
shape=[batch_size, sparse_seq_num, -1,
embedding_size]) #[B, seq_num, per_seq_len, e]
if params["parameters"]["combiner"] == "sum":
tmp_feat = tf.reduce_sum(deep_embedding_seq, axis=2)
else:
tmp_feat = tf.reduce_mean(deep_embedding_seq, axis=2)
sparse_seq_feat = tf.reshape(
tmp_feat,
[batch_size, sparse_seq_num * embedding_size]) #[B, seq_num*e]
else:
sparse_seq_feats = []
sparse_ids = []
for sparse_seq_name in sparse_seq_feat_list:
sp_ids = features[sparse_seq_name]
if params["args"].get("zero_padding2"):
sparse_seq_unique_ids, sparse_seq_unique_idx, _ = tf.unique_with_counts(
tf.reshape(sp_ids, [-1]))
deep_sparse_seq_weights = tf.reshape(
dynamic_embedding.embedding_lookup(
params=deep_dynamic_variables,
ids=sparse_seq_unique_ids,
name="deep_sparse_weights_{}".format(
sparse_seq_name)), [-1, embedding_size])
deep_embed_seq = tf.where(
tf.not_equal(
tf.expand_dims(sparse_seq_unique_ids, axis=1),
tf.zeros_like(
tf.expand_dims(sparse_seq_unique_ids,
axis=1))),
deep_sparse_seq_weights,
tf.zeros_like(deep_sparse_seq_weights))
deep_embedding_seq = tf.reshape(
tf.gather(deep_embed_seq, sparse_seq_unique_idx),
shape=[batch_size, -1, embedding_size])
if params["parameters"]["combiner"] == "sum":
tmp_feat = tf.reduce_sum(deep_embedding_seq, axis=1)
else:
tmp_feat = tf.reduce_mean(deep_embedding_seq, axis=1)
sparse_ids.append(sparse_seq_unique_ids)
sparse_seq_feats.append(
tf.reshape(tmp_feat, [batch_size, embedding_size]))
else:
tmp_feat = dynamic_embedding.safe_embedding_lookup_sparse(
embedding_weights=deep_dynamic_variables,
sparse_ids=sp_ids,
combiner=params["parameters"]["combiner"],
name="safe_embedding_lookup_sparse")
temp_uni_id, _, _ = tf.unique_with_counts(
tf.reshape(sp_ids.values, [-1]))
sparse_ids.append(temp_uni_id)
sparse_seq_feats.append(
tf.reshape(tmp_feat, [batch_size, embedding_size]))
sparse_seq_feat = tf.concat(sparse_seq_feats, axis=1)
sparse_seq_unique_ids, _ = tf.unique(tf.concat(sparse_ids, axis=0))
dense_feat = None
if dense_feat_num > 0:
logging.info("------ build dense feature -------")
den_id_list = sorted(dense_feat_list)
dense_feat_base = tf.concat(
[tf.reshape(features[str(i)], [-1, 1]) for i in den_id_list],
axis=1)
#deep_dense_w1 = tf.compat.v1.get_variable('deep_dense_w1',
# tf.TensorShape([dense_feat_num]),
# initializer=tf.compat.v1.truncated_normal_initializer(
# 2.0 / math.sqrt(dense_feat_num)),
# dtype=tf.float32)
#deep_dense_w2 = tf.compat.v1.get_variable('deep_dense_w2',
# tf.TensorShape([dense_feat_num]),
# initializer=tf.compat.v1.truncated_normal_initializer(
# 2.0 / math.sqrt(dense_feat_num)),
# dtype=tf.float32)
#w1 = tf.tile(tf.expand_dims(deep_dense_w1, axis=0), [tf.shape(dense_feat_base)[0], 1])
#dense_input_1 = tf.multiply(dense_feat_base, w1)
#dense_feat = dense_input_1
dense_feat = dense_feat_base
dense_seq_feat = None
if dense_seq_feat_num > 0:
logging.info("------ build dense seq feature -------")
den_seq_id_list = sorted(dense_seq_feat_list)
dense_seq_feat = tf.concat([
tf.reshape(features[str(i[0])], [-1, i[1]])
for i in den_seq_id_list
],
axis=1)
logging.info("------ join all feature -------")
fc_inputs = tf.concat([
x for x in [sparse_feat, sparse_seq_feat, dense_feat, dense_seq_feat]
if x is not None
],
axis=1)
logging.info("---- tracy debug input is ----")
logging.info(sparse_feat)
logging.info(sparse_seq_feat)
logging.info(dense_feat)
logging.info(dense_seq_feat)
logging.info(fc_inputs)
logging.info("------ join fc -------")
for idx, units in enumerate(params["parameters"]["hidden_units"]):
fc_inputs = fully_connected_with_bn_ahead(
inputs=fc_inputs,
num_outputs=units,
l2_reg=l2_reg,
scope="out_mlp_{}".format(idx),
activation_fn=tf.nn.relu,
train_phase=is_training,
use_bn=use_bn)
y_deep_ctr = fully_connected_with_bn_ahead(inputs=fc_inputs,
num_outputs=1,
activation_fn=tf.identity,
l2_reg=l2_reg,
scope="ctr_mlp",
train_phase=is_training,
use_bn=use_bn)
logging.info("------ build ctr out -------")
sample_rate = params["args"]["sample_rate"]
logit = tf.reshape(y_deep_ctr, shape=[-1], name="logit")
sample_logit = get_sample_logits(logit, sample_rate)
pred_ctr = tf.nn.sigmoid(logit, name="pred_ctr")
sample_pred_ctr = tf.nn.sigmoid(sample_logit, name="sample_pred_ctr")
logging.info("------ build predictions -------")
preds = {
'p_ctr': tf.reshape(pred_ctr, shape=[-1, 1]),
}
logging.info("---- deep_dynamic_variables.size ----")
logging.info(deep_dynamic_variables.size())
size = tf.identity(deep_dynamic_variables.size(), name="size")
label_col = "is_clk"
if params["args"].get("set_train_labels"):
label_col = params["args"]["set_train_labels"]["1"]
logging.info(
"------ build labels, label_col: {} -------".format(label_col))
if has_label:
labels_ctr = tf.reshape(features["is_clk"],
shape=[-1],
name="labels_ctr")
if mode == tf.estimator.ModeKeys.PREDICT:
logging.info("---- build tf-serving predict ----")
pred_cvr = tf.fill(tf.shape(pred_ctr), 1.0)
preds.update({
'labels_cart': tf.reshape(pred_cvr, shape=[-1, 1]),
'p_car': tf.reshape(features["dense_1608"], shape=[-1, 1]),
'labels_cvr': tf.reshape(pred_cvr, shape=[-1, 1]),
'p_cvr': tf.reshape(pred_cvr, shape=[-1, 1]),
})
if 'logid' in features:
preds.update(
{'logid': tf.reshape(features["logid"], shape=[-1, 1])})
if has_label:
logging.info("------ build offline label -------")
preds["labels_ctr"] = tf.reshape(labels_ctr, shape=[-1, 1])
export_outputs = {
"predict_export_outputs":
tf.estimator.export.PredictOutput(outputs=preds)
}
return tf.estimator.EstimatorSpec(mode,
predictions=preds,
export_outputs=export_outputs)
logging.info("----all vars:-----" + str(tf.compat.v1.global_variables()))
for var in tf.compat.v1.trainable_variables():
logging.info("----trainable------" + str(var))
logging.info("------ build metric -------")
#loss = tf.reduce_mean(
# tf.compat.v1.losses.log_loss(labels=labels_ctr, predictions=sample_pred_ctr),
# name="loss")
loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels_ctr, logits=sample_logit),
name="loss")
ctr_auc = tf.compat.v1.metrics.auc(labels=labels_ctr,
predictions=sample_pred_ctr,
name="ctr_auc")
label_ctr_avg = tf.reduce_mean(labels_ctr, name="label_ctr_avg")
real_pred_ctr_avg = tf.reduce_mean(pred_ctr, name="real_pred_ctr_avg")
sample_pred_ctr_avg = tf.reduce_mean(sample_pred_ctr, name="pred_ctr_avg")
sample_pred_bias_avg = tf.add(sample_pred_ctr_avg,
tf.negative(label_ctr_avg),
name="pred_bias_avg")
tf.compat.v1.summary.histogram('labels_ctr', labels_ctr)
tf.compat.v1.summary.histogram('pred_ctr', sample_pred_ctr)
tf.compat.v1.summary.histogram('real_pred_ctr', pred_ctr)
tf.compat.v1.summary.scalar('label_ctr_avg', label_ctr_avg)
tf.compat.v1.summary.scalar('pred_ctr_avg', sample_pred_ctr_avg)
tf.compat.v1.summary.scalar('real_pred_ctr_avg', real_pred_ctr_avg)
tf.compat.v1.summary.scalar('pred_bias_avg', sample_pred_bias_avg)
tf.compat.v1.summary.scalar('loss', loss)
tf.compat.v1.summary.scalar('ctr_auc', ctr_auc[1])
logging.info("------ compute l2 reg -------")
if params["parameters"]["use_l2"]:
all_unique_ids, _ = tf.unique(
tf.concat([
x for x in [sparse_unique_ids, sparse_seq_unique_ids]
if x is not None
],
axis=0))
all_unique_ids_w = dynamic_embedding.embedding_lookup(
deep_dynamic_variables,
all_unique_ids,
name="unique_ids_weights",
return_trainable=False)
embed_loss = l2_reg * tf.nn.l2_loss(
tf.reshape(all_unique_ids_w,
shape=[-1, embedding_size])) + tf.reduce_sum(
tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES))
tf.compat.v1.summary.scalar('embed_loss', embed_loss)
loss = loss + embed_loss
loss = tf.identity(loss, name="total_loss")
tf.compat.v1.summary.scalar('total_loss', loss)
if mode == tf.estimator.ModeKeys.EVAL:
logging.info("------ EVAL -------")
eval_metric_ops = {
"ctr_auc_eval": ctr_auc,
}
if has_label:
logging.info("------ build offline label -------")
preds["labels_ctr"] = tf.reshape(labels_ctr, shape=[-1, 1])
export_outputs = {
"predict_export_outputs":
tf.estimator.export.PredictOutput(outputs=preds)
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
export_outputs=export_outputs)
logging.info("---- Learning rate ----")
lr = get_learning_rate(params["parameters"]["learning_rate"],
params["parameters"]["use_decay"])
if mode == tf.estimator.ModeKeys.TRAIN:
global_step = tf.compat.v1.train.get_global_step()
logging.info("------ TRAIN -------")
optimizer_type = params["parameters"].get('optimizer', 'Adam')
if optimizer_type == 'Sgd':
optimizer = tf.compat.v1.train.GradientDescentOptimizer(
learning_rate=lr)
elif optimizer_type == 'Adagrad':
optimizer = tf.compat.v1.train.AdagradOptimizer(learning_rate=lr)
elif optimizer_type == 'Rmsprop':
optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=lr)
elif optimizer_type == 'Ftrl':
optimizer = tf.compat.v1.train.FtrlOptimizer(learning_rate=lr)
elif optimizer_type == 'Momentum':
optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=lr,
momentum=0.9)
else:
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=lr,
beta1=0.9,
beta2=0.999,
epsilon=1e-8)
if params["args"].get("addon_embedding"):
optimizer = dynamic_embedding.DynamicEmbeddingOptimizer(optimizer)
train_op = optimizer.minimize(loss, global_step=global_step)
# fix tf2 batch_normalization bug
update_ops = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.UPDATE_OPS)
logging.info('train ops: {}, update ops: {}'.format(
str(train_op), str(update_ops)))
train_op = tf.group([train_op, update_ops])
return tf.estimator.EstimatorSpec(mode=mode,
predictions=preds,
loss=loss,
train_op=train_op)
def fn():
embedding_weights = _random_weights(num_shards=3,
key_dtype=dtypes.int32)
sparse_ids, sparse_weights = ids_and_weights_2d()
de.safe_embedding_lookup_sparse(embedding_weights, sparse_ids,
sparse_weights)
def test_safe_embedding_lookup_sparse_with_initializer(self):
id = 0
embed_dim = 8
dense_shape = np.array([64, 128, 32])
total_space = 64 * 128 * 32
elements_num = int(total_space * 0.50)
for initializer, target_mean, target_stddev in [
(init_ops.random_normal_initializer(0.0, 0.001), 0.0, 0.00029),
(init_ops.truncated_normal_initializer(0.0, 0.001), 0.0, 0.00029),
(keras_init_ops.RandomNormalV2(mean=0.0, stddev=0.001), 0.0, 0.00029),
]:
with self.session(config=default_config,
use_gpu=test_util.is_gpu_available()):
id += 1
embedding_weights = de.get_variable(
"safe-init-bugfix-" + str(id),
key_dtype=dtypes.int64,
value_dtype=dtypes.float32,
devices=_get_devices() * 3,
initializer=initializer,
dim=embed_dim,
)
indices_1d = np.random.randint(0, total_space, elements_num)
indices_1d = np.unique(indices_1d)
indices_1d.sort()
indices_3d = []
for _i in range(indices_1d.size):
a_indice = []
quotient = int(indices_1d[_i] / (128 * 32))
remainder = indices_1d[_i] % (128 * 32)
a_indice.append(quotient)
quotient = int(remainder / 32)
remainder = remainder % 32
a_indice.extend([quotient, remainder])
indices_3d.extend([a_indice])
indices_3d = np.array(indices_3d)
ids = np.random.randint(
-0x7FFFFFFFFFFFFFFF,
0x7FFFFFFFFFFFFFFF,
indices_1d.size,
dtype=np.int64,
)
sparse_ids = sparse_tensor.SparseTensor(
constant_op.constant(indices_3d, dtypes.int64),
constant_op.constant(ids, dtypes.int64),
constant_op.constant(dense_shape, dtypes.int64),
)
vals_op = de.safe_embedding_lookup_sparse(embedding_weights,
sparse_ids,
None,
combiner="mean").eval()
mean = self.evaluate(math_ops.reduce_mean(vals_op))
stddev = self.evaluate(math_ops.reduce_std(vals_op))
rtol = 2e-4
atol = rtol
self.assertTrue(not (vals_op[0][0][0] == vals_op[0][0][1]))
self.assertAllClose(target_mean, mean, rtol, atol)
self.assertAllClose(target_stddev, stddev, 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(
allow_soft_placement=True,
gpu_options=config_pb2.GPUOptions(allow_growth=True),
)
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],
):
with self.session(config=config,
use_gpu=test_util.is_gpu_available()):
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],
)
base_embedding = embedding_ops.safe_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 var prepare
embeddings = de.get_variable(
"s6030-" + 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 branch
test_var, trainable = de.safe_embedding_lookup_sparse(
embeddings,
sp_ids,
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),
)
