def metric_fn(logits, dup_mask, params):
dup_mask = tf.cast(dup_mask, tf.float32)
logits = tf.slice(logits, [0, 1], [-1, -1])
in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
logits, dup_mask, params["match_mlperf"])
metric_weights = tf.cast(metric_weights, tf.float32)
return in_top_k, metric_weights
def tpu_step_fn(total, count, duplicate_mask):
"""One step in evaluation."""
logits = logits_fn(embedding, params)
in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
logits, duplicate_mask, FLAGS.ml_perf)
metric_weights = tf.cast(metric_weights, tf.float32)
total += tf.reduce_sum(tf.multiply(in_top_k, metric_weights))
count += tf.reduce_sum(metric_weights)
return total, count
def step_fn(inputs):
"""Computes eval metrics per replica."""
features, _ = inputs
softmax_logits = keras_model([features[movielens.USER_COLUMN],
features[movielens.ITEM_COLUMN]])
logits = tf.slice(softmax_logits, [0, 0, 1], [-1, -1, -1])
dup_mask = features[rconst.DUPLICATE_MASK]
in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
logits,
dup_mask,
params["match_mlperf"])
metric_weights = tf.cast(metric_weights, tf.float32)
hr_sum = tf.reduce_sum(in_top_k*metric_weights)
hr_count = tf.reduce_sum(metric_weights)
return hr_sum, hr_count
