def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self._parse_sequence_weight(features)
features = self.sparse2dense(features,
self._dataset.varlen_list)
duration, labels = labels
features = self._dense2sparse(features, self._dataset.varlen_list)
network = self._Network(self._flags, self._dataset, 'network')
dense, embeddings = network.build_features(features)
network_out = network(dense, embeddings, is_training)
predictions = tf.keras.layers.Dense(1,
activation=tf.sigmoid,
name='output')(network_out)
if mode == estimator.ModeKeys.PREDICT:
outputs = self._build_predict_outputs(features, predictions)
return estimator.EstimatorSpec(mode, predictions=outputs)
metrics = self._build_duration_metrics(duration, labels,
predictions)
weights = self._build_weights(duration, labels)
loss = self._build_loss(labels, predictions, weights=weights)
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode):
_ = labels
step = training.get_global_step()
w = variable_scope.get_variable(
'w',
shape=[],
initializer=init_ops.zeros_initializer(),
dtype=dtypes.int64)
if estimator_lib.ModeKeys.TRAIN == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
step_inc = state_ops.assign_add(training.get_global_step(),
1)
with ops.control_dependencies([step_inc]):
assign_w_to_step_plus_2 = w.assign(step + 2)
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
train_op=assign_w_to_step_plus_2)
if estimator_lib.ModeKeys.EVAL == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
loss = constant_op.constant(5.)
return estimator_lib.EstimatorSpec(
mode,
loss=loss,
# w is constant in each step, so the mean.
# w = 0 if step==0 else step+2
eval_metric_ops={'mean_of_const': metrics_lib.mean(w)})
def model_fn(features, labels, mode):
_ = labels
step = tf.compat.v1.train.get_global_step()
w = tf.compat.v1.get_variable(
'w',
shape=[],
initializer=tf.compat.v1.initializers.zeros(),
dtype=tf.dtypes.int64)
if estimator_lib.ModeKeys.TRAIN == mode:
# to consume features, we have control dependency
with tf.control_dependencies([features]):
step_inc = tf.compat.v1.assign_add(
tf.compat.v1.train.get_global_step(), 1)
with tf.control_dependencies([step_inc]):
assign_w_to_step_plus_2 = w.assign(step + 2)
return estimator_lib.EstimatorSpec(
mode,
loss=tf.constant(3.),
train_op=assign_w_to_step_plus_2)
if estimator_lib.ModeKeys.EVAL == mode:
# to consume features, we have control dependency
with tf.control_dependencies([features]):
loss = tf.constant(5.)
mean = metrics_module.Mean()
mean.update_state(w)
return estimator_lib.EstimatorSpec(
mode,
loss=loss,
# w is constant in each step, so the mean.
# w = 0 if step==0 else step+2
eval_metric_ops={'mean_of_const': mean})
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self._parse_sequence_weight(features)
features = self.sparse2dense(features,
self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
network = self._Network(self._flags, self._dataset, 'network')
dense, embeddings = network.build_features(features)
network_out = network(dense, embeddings, is_training)
# 防止回归值小于0
predictions = tf.maximum(
tf.keras.layers.Dense(1, activation=None,
name='output')(network_out), 0.)
if mode == estimator.ModeKeys.PREDICT:
outputs = {
"predictions": predictions,
self._flags.label_key: tf.expm1(predictions)
}
self._output_cols = list(outputs.keys())
return estimator.EstimatorSpec(mode, predictions=outputs)
loss = self._build_regression_loss(labels, predictions)
metrics = self._build_regression_metrics(loss, labels, predictions,
self._flags.label_key)
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self.sparse2dense(features, self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
self.right_features = self._get_features_by_index(features, 1)
predict_tower = self._Network(self._flags, self._dataset, 'predict')
p_dense, p_embeddings = predict_tower.build_features(self.right_features)
p_tower_out = predict_tower(p_dense, p_embeddings, is_training)
p_out = tf.keras.layers.Dense(1, activation=tf.sigmoid, name = 'p_out')(p_tower_out)
if mode == estimator.ModeKeys.PREDICT:
outputs = self._build_predict_outputs(features, p_out)
return estimator.EstimatorSpec(
mode,
predictions=outputs
)
self.left_features = self._get_features_by_index(features, 0)
bias_net = network_factory(self._flags.pal_submodel)
bias_tower = bias_net(self._flags, self._dataset, 'bias', hidden=self._flags.pal_bias)
b_dense, b_embeddings = bias_tower.build_features(self.left_features)
b_tower_out = bias_tower(b_dense, b_embeddings, is_training)
b_out = tf.keras.layers.Dense(1, activation=tf.sigmoid, name = 'b_out')(b_tower_out)
predictions = tf.multiply(b_out, p_out, name=self.predictions)
metrics = self._build_metrics(labels, predictions)
loss = self._build_loss(labels, predictions)
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def model_fn(features, labels, mode):
_ = labels
if estimator_lib.ModeKeys.TRAIN == mode:
with ops.control_dependencies([features]):
train_op = state_ops.assign_add(training.get_global_step(), 1)
return estimator_lib.EstimatorSpec(
mode, loss=constant_op.constant(3.), train_op=train_op)
if estimator_lib.ModeKeys.EVAL == mode:
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(5.),
eval_metric_ops={'mean_of_features': metrics_lib.mean(features)})
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self.sparse2dense(features,
self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
self.left_features = self._get_features_by_index(features, 0)
self.right_features = self._get_features_by_index(features, 1)
left_network = self._Network(self._flags, self._dataset, 'left')
right_network = self._Network(self._flags, self._dataset, 'right')
u_dense, u_embeddings = left_network.build_features(
self.left_features)
v_dense, v_embeddings = right_network.build_features(
self.right_features)
left_embedding = tf.reduce_mean(tf.concat(u_embeddings, 1), 1)
right_embedding = tf.reduce_mean(tf.concat(v_embeddings, 1), 1)
dense = u_dense + v_dense
embeddings = u_embeddings + v_embeddings
network_out = left_network(dense, embeddings, is_training)
predictions = tf.keras.layers.Dense(
1, activation=None, name='predictions')(network_out)
predictions = tf.identity(predictions, name=self.predictions)
if mode == estimator.ModeKeys.PREDICT:
left_bias = tf.reduce_mean(tf.ones_like(left_embedding),
-1,
keepdims=True)
left_embedding = tf.concat((left_embedding, left_bias),
-1,
name=self.left_output)
right_bias = tf.reduce_mean(right_network(
v_dense, v_embeddings, False),
-1,
keepdims=True)
right_embedding = tf.concat((right_embedding, right_bias),
-1,
name=self.right_output)
outputs = self._build_predict_outputs(features,
right_embedding)
return estimator.EstimatorSpec(mode, predictions=outputs)
loss = self._build_regression_loss(labels, predictions)
metrics = self._build_regression_metrics(loss, labels, predictions,
'duration')
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self.sparse2dense(features,
self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
# 网络
network = self._Network(self._flags, self._dataset,
'dnn') # 初始化网络对象
dense_feat, embedding_feat = network.build_features(
features) # 特征embedding化，[batch_size, length]
tf.logging.debug("feature dense: %r" % dense_feat)
tf.logging.debug("feature embed: %r" % embedding_feat)
network_out = network(dense_feat, embedding_feat,
is_training) # 模型输出
tf.logging.debug("network out: %r" % network_out.shape)
# user embedding
user_embedding = tf.keras.layers.Dense(
self._flags.vector_dim)(network_out)
tf.logging.debug("user embedding: %r" % user_embedding.shape)
if mode == estimator.ModeKeys.PREDICT:
predictions = {"user_embedding": user_embedding}
return estimator.EstimatorSpec(mode, predictions=predictions)
# softmax层
nce_weights = tf.Variable(tf.truncated_normal(
[self.item_num, self._flags.vector_dim],
stddev=1.0 / tf.math.sqrt(float(self._flags.vector_dim))),
name="nce_weights")
logits = tf.matmul(user_embedding, tf.transpose(
nce_weights)) # 计算所有item的概率，[batch_size, item_num]
labels = tf.reshape(
labels, [-1, self.label_num]) # [batch_size, label_num]
if mode == tf.estimator.ModeKeys.EVAL:
eval_loss = self.build_eval_loss(logits, labels)
metrics = self.build_metrics(logits, labels)
return tf.estimator.EstimatorSpec(mode,
loss=eval_loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
loss = self.build_loss(user_embedding, labels,
nce_weights) # 训练使用nce_loss
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self.sparse2dense(features,
self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
self.left_features = self._get_features_by_index(features, 0)
self.right_features = self._get_features_by_index(features, 1)
user_tower = self._Network(self._flags, self._dataset, 'user')
item_tower = self._Network(self._flags, self._dataset, 'item')
u_dense, u_embeddings = user_tower.build_features(
self.left_features)
v_dense, v_embeddings = item_tower.build_features(
self.right_features)
u_tower_out = user_tower(u_dense, u_embeddings, is_training)
v_tower_out = item_tower(v_dense, v_embeddings, is_training)
u_vector = tf.nn.l2_normalize(tf.keras.layers.Dense(
self._flags.vector_dim)(u_tower_out),
axis=-1,
name=self.left_output)
v_vector = tf.nn.l2_normalize(tf.keras.layers.Dense(
self._flags.vector_dim)(v_tower_out),
axis=-1,
name=self.right_output)
# 使用余弦相似度，计算结果在[-1, 1]，变换成[0, 1]
predictions = tf.divide(1. + tf.reduce_sum(
tf.multiply(u_vector, v_vector), axis=-1, keep_dims=True),
2.,
name=self.predictions)
if mode == estimator.ModeKeys.PREDICT:
if self._flags.predict_with_emb:
outputs = self._build_predict_outputs(features, v_vector)
elif self._flags.predict_with_user:
outputs = self._build_predict_outputs_user(
features, u_vector)
else:
outputs = super(DssmModel, self)._build_predict_outputs(
features, predictions)
return estimator.EstimatorSpec(mode, predictions=outputs)
metrics = self._build_metrics(labels, predictions)
loss = self._build_loss(labels, predictions)
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def _serving_ops(self, features):
"""Add ops for serving to the graph."""
with tf.compat.v1.variable_scope("model", use_resource=True):
filtering_features = {}
prediction_features = {}
values_length = tf.compat.v1.shape(
features[feature_keys.FilteringFeatures.VALUES])[1]
for key, value in features.items():
if key == feature_keys.State.STATE_TUPLE:
# Ignore state input. The model's default start state is replicated
# across the batch.
continue
if key == feature_keys.FilteringFeatures.VALUES:
filtering_features[key] = value
else:
filtering_features[key] = value[:, :values_length]
prediction_features[key] = value[:, values_length:]
cold_filtering_outputs = self.model.define_loss(
features=filtering_features, mode=estimator_lib.ModeKeys.EVAL)
prediction_features[feature_keys.State.STATE_TUPLE] = (
cold_filtering_outputs.end_state)
with tf.compat.v1.variable_scope("model", reuse=True):
prediction_outputs = self.model.predict(
features=prediction_features)
return estimator_lib.EstimatorSpec(
mode=estimator_lib.ModeKeys.PREDICT,
export_outputs={
feature_keys.SavedModelLabels.PREDICT:
_NoStatePredictOutput(prediction_outputs),
},
# Likely unused, but it is necessary to return `predictions` to satisfy
# the Estimator's error checking.
predictions={})
def _serving_ops(self, features):
"""Add ops for serving to the graph."""
with tf.compat.v1.variable_scope("model", use_resource=True):
prediction_outputs = self.model.predict(features=features)
with tf.compat.v1.variable_scope("model", reuse=True):
filtering_outputs = self.create_loss(features,
estimator_lib.ModeKeys.EVAL)
with tf.compat.v1.variable_scope("model", reuse=True):
no_state_features = {
k: v
for k, v in features.items()
if not k.startswith(feature_keys.State.STATE_PREFIX)
}
# Ignore any state management when cold-starting. The model's default
# start state is replicated across the batch.
cold_filtering_outputs = self.model.define_loss(
features=no_state_features, mode=estimator_lib.ModeKeys.EVAL)
return estimator_lib.EstimatorSpec(
mode=estimator_lib.ModeKeys.PREDICT,
export_outputs={
feature_keys.SavedModelLabels.PREDICT:
export_lib.PredictOutput(prediction_outputs),
feature_keys.SavedModelLabels.FILTER:
export_lib.PredictOutput(
state_to_dictionary(filtering_outputs.end_state)),
feature_keys.SavedModelLabels.COLD_START_FILTER:
_NoStatePredictOutput(
state_to_dictionary(cold_filtering_outputs.end_state))
},
# Likely unused, but it is necessary to return `predictions` to satisfy
# the Estimator's error checking.
predictions={})
def _evaluate_ops(self, features):
"""Add ops for evaluation (aka filtering) to the graph."""
mode = estimator_lib.ModeKeys.EVAL
with tf.compat.v1.variable_scope("model", use_resource=True):
model_outputs = self.create_loss(features, mode)
metrics = {}
# Just output in-sample predictions for the last chunk seen
for prediction_key, prediction_value in model_outputs.predictions.items(
):
metrics[prediction_key] = _identity_metric_single(
prediction_key, prediction_value)
metrics[feature_keys.FilteringResults.TIMES] = _identity_metric_single(
feature_keys.FilteringResults.TIMES,
model_outputs.prediction_times)
metrics[feature_keys.FilteringResults.STATE_TUPLE] = (
_identity_metric_nested(feature_keys.FilteringResults.STATE_TUPLE,
model_outputs.end_state))
metrics[metric_keys.MetricKeys.LOSS_MEAN] = tf.compat.v1.metrics.mean(
model_outputs.loss, name="average_loss")
return estimator_lib.EstimatorSpec(
loss=model_outputs.loss,
mode=mode,
eval_metric_ops=metrics,
# needed for custom metrics.
predictions=model_outputs.predictions)
def _predict_ops(self, features):
"""Add ops for prediction to the graph."""
with tf.compat.v1.variable_scope("model", use_resource=True):
prediction = self.model.predict(features=features)
prediction[feature_keys.PredictionResults.TIMES] = features[
feature_keys.PredictionFeatures.TIMES]
return estimator_lib.EstimatorSpec(predictions=prediction,
mode=estimator_lib.ModeKeys.PREDICT)
def model_fn(features, mode):
del features
global_step = training.get_global_step()
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant([5.]),
predictions={'x': constant_op.constant([5.])},
train_op=global_step.assign_add(1))
def model_fn(features, labels, mode):
_ = labels
if estimator_lib.ModeKeys.TRAIN == mode:
with tf.control_dependencies([features]):
train_op = tf.compat.v1.assign_add(
tf.compat.v1.train.get_global_step(), 1)
return estimator_lib.EstimatorSpec(mode,
loss=tf.constant(3.),
train_op=train_op)
if estimator_lib.ModeKeys.EVAL == mode:
mean = metrics_module.Mean()
mean.update_state(features)
return estimator_lib.EstimatorSpec(mode,
loss=tf.constant(5.),
eval_metric_ops={
'mean_of_features':
mean,
})
def model_fn(features, labels, mode):
_, _ = features, labels
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(saver=training.Saver()),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': metrics_lib.mean(constant_op.constant(2.))
})
def model_fn(features, labels, mode):
_, _ = features, labels
mean = metrics_module.Mean()
mean.update_state(constant_op.constant(2.))
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(saver=training.Saver()),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': mean,
})
def _train_ops(self, features):
"""Add training ops to the graph."""
mode = estimator_lib.ModeKeys.TRAIN
with tf.compat.v1.variable_scope(
"model",
# Use ResourceVariables to avoid race conditions.
use_resource=True):
model_outputs = self.create_loss(features, mode)
train_op = self.optimizer.minimize(
model_outputs.loss, global_step=tf.compat.v1.train.get_global_step())
return estimator_lib.EstimatorSpec(
loss=model_outputs.loss, mode=mode, train_op=train_op)
def model_fn(features, labels, mode):
_, _ = features, labels
mean = tf.keras.metrics.Mean()
mean.update_state(tf.constant(2.))
return estimator_lib.EstimatorSpec(
mode,
loss=tf.constant(3.),
scaffold=tf.compat.v1.train.Scaffold(
saver=tf.compat.v1.train.Saver()),
train_op=tf.constant(5.),
eval_metric_ops={
'mean_of_features': mean,
})
def model_fn(features, labels, mode):
_, _ = features, labels
def init_fn(scaffold, session):
_, _ = scaffold, session
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(init_fn=init_fn),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': metrics_lib.mean(constant_op.constant(2.))
})
def model_fn(features, labels, mode):
_, _ = features, labels
def init_fn(scaffold, session):
_, _ = scaffold, session
mean = metrics_module.Mean()
mean.update_state(tf.constant(2.))
return estimator_lib.EstimatorSpec(
mode,
loss=tf.constant(3.),
scaffold=tf.compat.v1.train.Scaffold(init_fn=init_fn),
train_op=tf.constant(5.),
eval_metric_ops={
'mean_of_features': mean,
})
def model_fn(features, labels, mode):
_, _ = features, labels
w = variables.VariableV1(
initial_value=[0.],
trainable=False,
collections=[ops.GraphKeys.SAVEABLE_OBJECTS])
init_op = control_flow_ops.group(
[w.initializer, training.get_global_step().initializer])
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(init_op=init_op),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': metrics_lib.mean(constant_op.constant(2.))
})
def model_fn(features, labels, mode):
_, _ = features, labels
w = tf.compat.v1.Variable(
initial_value=[0.],
trainable=False,
collections=[tf.compat.v1.GraphKeys.SAVEABLE_OBJECTS])
init_op = tf.group([
w.initializer,
tf.compat.v1.train.get_global_step().initializer
])
mean = metrics_module.Mean()
mean.update_state(tf.constant(2.))
return estimator_lib.EstimatorSpec(
mode,
loss=tf.constant(3.),
scaffold=tf.compat.v1.train.Scaffold(init_op=init_op),
train_op=tf.constant(5.),
eval_metric_ops={
'mean_of_features': mean,
})
def model_fn(features, labels, mode, params, config):
is_training = True if mode == estimator.ModeKeys.TRAIN else False
if mode != estimator.ModeKeys.PREDICT:
features = self.sparse2dense(features,
self._dataset.varlen_list)
features = self._dense2sparse(features, self._dataset.varlen_list)
# 左塔
left_tower = self._Network(self._flags, self._dataset,
'left_tower') # 初始化网络对象
self.left_features = self._get_features_by_index(features, 0)
left_dense, left_embeddings = left_tower.build_features(
self.left_features
) # 特征embedding化，[batch_size * pre_batch, length]
tf.logging.debug("LeftTower input dense: %r" % left_dense)
tf.logging.debug("LeftTower input embed: %r" % left_embeddings)
left_tower_out = left_tower(left_dense, left_embeddings,
is_training) # 输出
left_tower_embedding = tf.nn.l2_normalize(
tf.keras.layers.Dense(self._flags.vector_dim)(left_tower_out),
axis=-1,
name=self.left_output
) # 单位化，[batch_size * pre_batch, embedding_size]
tf.logging.debug("LeftTower output: %r" %
left_tower_embedding.shape)
# 右塔
right_tower = self._Network(self._flags, self._dataset,
'right_tower') # 初始化网络对象
self.right_features = self._get_features_by_index(features, 1)
right_dense, right_embeddings = right_tower.build_features(
self.right_features
) # 特征embedding化，[batch_size * pre_batch * sample_num, length]
tf.logging.debug("RightTower input dense: %r" % right_dense)
tf.logging.debug("RightTower input embed: %r" % right_embeddings)
right_tower_out = right_tower(right_dense, right_embeddings,
is_training) # 输出
right_tower_embedding = tf.nn.l2_normalize(
tf.keras.layers.Dense(self._flags.vector_dim)(right_tower_out),
axis=-1,
name=self.right_output
) # 单位化，[batch_size * pre_batch * sample_num, embedding_size]
tf.logging.debug("RightTower output: %r" %
right_tower_embedding.shape)
if mode == estimator.ModeKeys.PREDICT:
outputs = self._build_predict_outputs(features,
right_tower_embedding)
return estimator.EstimatorSpec(mode, predictions=outputs)
# 损失
if self.sample_num > 1: # 如果是多个sample，那么需要做reshape，使用softmax做分类
left_tower_embedding = tf.expand_dims(
left_tower_embedding,
1) # [batch_size * pre_batch, 1, embedding_size]
tf.logging.debug("LeftTower embedding: %r" %
left_tower_embedding.shape)
right_tower_embedding = tf.reshape(
right_tower_embedding,
[-1, self.sample_num, right_tower_embedding.shape[-1]
]) # [batch_size * pre_batch, sample_num, embedding_size]
tf.logging.debug("RightTower embedding: %r" %
right_tower_embedding.shape)
cosine_score = tf.matmul(
left_tower_embedding,
right_tower_embedding,
transpose_b=True
) # [batch_size * pre_batch, 1, sample_num]
cosine_score = tf.squeeze(
cosine_score, axis=1, name=self.predictions
) # [batch_size * pre_batch, sample_num]
tf.logging.debug("Cosine: %r" % cosine_score.shape)
labels = tf.reshape(labels,
[-1, self.sample_num
]) # [batch_size * pre_batch, sample_num]
tf.logging.debug("Labels: %r" % labels.shape)
loss = self.build_loss(labels, cosine_score)
metrics = self.build_softmax_metric(labels, cosine_score)
else: # 如果是单个sample，那么使用logloss
cosine_score = tf.divide(1. + tf.reduce_sum(tf.multiply(
left_tower_embedding, right_tower_embedding),
axis=-1,
keep_dims=True),
2.,
name=self.predictions)
loss = self._build_loss(labels, cosine_score)
metrics = self._build_metrics(labels, cosine_score)
self._build_summary(loss, metrics)
if mode == estimator.ModeKeys.EVAL:
return estimator.EstimatorSpec(mode,
loss=loss,
eval_metric_ops=metrics)
assert mode == estimator.ModeKeys.TRAIN
train_op = self._build_train_op(loss)
return estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
