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