easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/core/learning_schedules.py

@@ -0,0 +1,228 @@
+# -*- encoding:utf-8 -*-
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Library of common learning rate schedules."""
+
+import numpy as np
+import tensorflow as tf
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+def exponential_decay_with_burnin(global_step,
+                                  learning_rate_base,
+                                  learning_rate_decay_steps,
+                                  learning_rate_decay_factor,
+                                  burnin_learning_rate=0.0,
+                                  burnin_steps=0,
+                                  min_learning_rate=0.0,
+                                  staircase=True):
+  """Exponential decay schedule with burn-in period.
+
+  In this schedule, learning rate is fixed at burnin_learning_rate
+  for a fixed period, before transitioning to a regular exponential
+  decay schedule.
+
+  Args:
+    global_step: int tensor representing global step.
+    learning_rate_base: base learning rate.
+    learning_rate_decay_steps: steps to take between decaying the learning rate.
+      Note that this includes the number of burn-in steps.
+    learning_rate_decay_factor: multiplicative factor by which to decay
+      learning rate.
+    burnin_learning_rate: initial learning rate during burn-in period.  If
+      0.0 (which is the default), then the burn-in learning rate is simply
+      set to learning_rate_base.
+    burnin_steps: number of steps to use burnin learning rate.
+    min_learning_rate: the minimum learning rate.
+    staircase: whether use staircase decay.
+
+  Returns:
+    a (scalar) float tensor representing learning rate
+  """
+  if burnin_learning_rate == 0:
+    burnin_rate = learning_rate_base
+  else:
+    slope = (learning_rate_base - burnin_learning_rate) / burnin_steps
+    burnin_rate = slope * tf.cast(global_step,
+                                  tf.float32) + burnin_learning_rate
+  post_burnin_learning_rate = tf.train.exponential_decay(
+      learning_rate_base,
+      global_step - burnin_steps,
+      learning_rate_decay_steps,
+      learning_rate_decay_factor,
+      staircase=staircase)
+  return tf.maximum(
+      tf.where(
+          tf.less(tf.cast(global_step, tf.int32), tf.constant(burnin_steps)),
+          burnin_rate, post_burnin_learning_rate),
+      min_learning_rate,
+      name='learning_rate')
+
+
+def cosine_decay_with_warmup(global_step,
+                             learning_rate_base,
+                             total_steps,
+                             warmup_learning_rate=0.0,
+                             warmup_steps=0,
+                             hold_base_rate_steps=0):
+  """Cosine decay schedule with warm up period.
+
+  Cosine annealing learning rate as described in:
+    Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts.
+    ICLR 2017. https://arxiv.org/abs/1608.03983
+  In this schedule, the learning rate grows linearly from warmup_learning_rate
+  to learning_rate_base for warmup_steps, then transitions to a cosine decay
+  schedule.
+
+  Args:
+    global_step: int64 (scalar) tensor representing global step.
+    learning_rate_base: base learning rate.
+    total_steps: total number of training steps.
+    warmup_learning_rate: initial learning rate for warm up.
+    warmup_steps: number of warmup steps.
+    hold_base_rate_steps: Optional number of steps to hold base learning rate
+      before decaying.
+
+  Returns:
+    a (scalar) float tensor representing learning rate.
+
+  Raises:
+    ValueError: if warmup_learning_rate is larger than learning_rate_base,
+      or if warmup_steps is larger than total_steps.
+  """
+  if learning_rate_base < warmup_learning_rate:
+    raise ValueError('learning_rate_base must be larger '
+                     'or equal to warmup_learning_rate.')
+  if total_steps < warmup_steps:
+    raise ValueError('total_steps must be larger or equal to ' 'warmup_steps.')
+  learning_rate = 0.5 * learning_rate_base * (1 + tf.cos(
+      np.pi *
+      (tf.cast(global_step, tf.float32) - warmup_steps - hold_base_rate_steps) /
+      float(total_steps - warmup_steps - hold_base_rate_steps)))
+  if hold_base_rate_steps > 0:
+    learning_rate = tf.where(global_step > warmup_steps + hold_base_rate_steps,
+                             learning_rate, learning_rate_base)
+  if warmup_steps > 0:
+    slope = (learning_rate_base - warmup_learning_rate) / warmup_steps
+    warmup_rate = slope * tf.cast(global_step,
+                                  tf.float32) + warmup_learning_rate
+    learning_rate = tf.where(global_step < warmup_steps, warmup_rate,
+                             learning_rate)
+  return tf.where(
+      global_step > total_steps, 0.0, learning_rate, name='learning_rate')
+
+
+def manual_stepping(global_step, boundaries, rates, warmup=False):
+  """Manually stepped learning rate schedule.
+
+  This function provides fine grained control over learning rates.  One must
+  specify a sequence of learning rates as well as a set of integer steps
+  at which the current learning rate must transition to the next.  For example,
+  if boundaries = [5, 10] and rates = [.1, .01, .001], then the learning
+  rate returned by this function is .1 for global_step=0,...,4, .01 for
+  global_step=5...9, and .001 for global_step=10 and onward.
+
+  Args:
+    global_step: int64 (scalar) tensor representing global step.
+    boundaries: a list of global steps at which to switch learning
+      rates.  This list is assumed to consist of increasing positive integers.
+    rates: a list of (float) learning rates corresponding to intervals between
+      the boundaries.  The length of this list must be exactly
+      len(boundaries) + 1.
+    warmup: Whether to linearly interpolate learning rate for steps in
+      [0, boundaries[0]].
+
+  Returns:
+    a (scalar) float tensor representing learning rate
+  Raises:
+    ValueError: if one of the following checks fails:
+      1. boundaries is a strictly increasing list of positive integers
+      2. len(rates) == len(boundaries) + 1
+      3. boundaries[0] != 0
+  """
+  if any([b < 0 for b in boundaries]) or any(
+      [not isinstance(b, int) for b in boundaries]):
+    raise ValueError('boundaries must be a list of positive integers')
+  if any([bnext <= b for bnext, b in zip(boundaries[1:], boundaries[:-1])]):
+    raise ValueError('Entries in boundaries must be strictly increasing.')
+  if any([not isinstance(r, float) for r in rates]):
+    raise ValueError('Learning rates must be floats')
+  if len(rates) != len(boundaries) + 1:
+    raise ValueError('Number of provided learning rates must exceed '
+                     'number of boundary points by exactly 1.')
+
+  if boundaries and boundaries[0] == 0:
+    raise ValueError('First step cannot be zero.')
+
+  if warmup and boundaries:
+    slope = (rates[1] - rates[0]) * 1.0 / boundaries[0]
+    warmup_steps = list(range(boundaries[0]))
+    warmup_rates = [rates[0] + slope * step for step in warmup_steps]
+    boundaries = warmup_steps + boundaries
+    rates = warmup_rates + rates[1:]
+  else:
+    boundaries = [0] + boundaries
+  num_boundaries = len(boundaries)
+  rate_index = tf.reduce_max(
+      tf.where(
+          tf.greater_equal(global_step, boundaries),
+          list(range(num_boundaries)), [0] * num_boundaries))
+  return tf.reduce_sum(
+      rates * tf.one_hot(rate_index, depth=num_boundaries),
+      name='learning_rate')
+
+
+def transformer_policy(global_step,
+                       learning_rate,
+                       d_model,
+                       warmup_steps,
+                       step_scaling_rate=1.0,
+                       max_lr=None,
+                       coefficient=1.0,
+                       dtype=tf.float32):
+  """Transformer's learning rate schedule.
+
+  Transformer's learning rate policy from
+  https://arxiv.org/pdf/1706.03762.pdf
+  with a hat (max_lr) (also called "noam" learning rate decay scheme).
+
+  Args:
+    global_step: global step TensorFlow tensor (ignored for this policy).
+    learning_rate (float): initial learning rate to use.
+    d_model (int): model dimensionality.
+    warmup_steps (int): number of warm-up steps.
+    step_scaling_rate (float): num step scale rate
+    max_lr (float): maximal learning rate, i.e. hat.
+    coefficient (float): optimizer adjustment.
+        Recommended 0.002 if using "Adam" else 1.0.
+    dtype: dtype for this policy.
+
+  Returns:
+    learning rate at step ``global_step``.
+  """
+  step_num = tf.cast(global_step, dtype=dtype)
+  ws = tf.cast(warmup_steps, dtype=dtype)
+  step_num *= step_scaling_rate
+  ws *= step_scaling_rate
+
+  decay = coefficient * d_model**-0.5 * tf.minimum((step_num + 1) * ws**-1.5,
+                                                   (step_num + 1)**-0.5)
+
+  new_lr = decay * learning_rate
+  if max_lr is not None:
+    return tf.minimum(max_lr, new_lr)
+  return new_lr

easy_rec/python/core/metrics.py

@@ -0,0 +1,402 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import json
+import logging
+import os
+from collections import defaultdict
+
+import numpy as np
+import tensorflow as tf
+from sklearn import metrics as sklearn_metrics
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+
+from easy_rec.python.utils.estimator_utils import get_task_index_and_num
+from easy_rec.python.utils.io_util import read_data_from_json_path
+from easy_rec.python.utils.io_util import save_data_to_json_path
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+def max_f1(label, predictions):
+  """Calculate the largest F1 metric under different thresholds.
+
+  Args:
+    label: Ground truth (correct) target values.
+    predictions: Estimated targets as returned by a model.
+  """
+  from easy_rec.python.core.easyrec_metrics import metrics_tf
+  num_thresholds = 200
+  kepsilon = 1e-7
+  thresholds = [
+      (i + 1) * 1.0 / (num_thresholds - 1) for i in range(num_thresholds - 2)
+  ]
+  thresholds = [0.0 - kepsilon] + thresholds + [1.0 + kepsilon]
+
+  f1_scores = []
+  precision_update_ops = []
+  recall_update_ops = []
+  for threshold in thresholds:
+    pred = (predictions > threshold)
+    precision, precision_update_op = metrics_tf.precision(
+        labels=label, predictions=pred, name='precision_%s' % threshold)
+    recall, recall_update_op = metrics_tf.recall(
+        labels=label, predictions=pred, name='recall_%s' % threshold)
+    f1_score = (2 * precision * recall) / (precision + recall + 1e-12)
+    precision_update_ops.append(precision_update_op)
+    recall_update_ops.append(recall_update_op)
+    f1_scores.append(f1_score)
+
+  f1 = tf.math.reduce_max(tf.stack(f1_scores))
+  f1_update_op = tf.group(precision_update_ops + recall_update_ops)
+  return f1, f1_update_op
+
+
+def _separated_auc_impl(labels, predictions, keys, reduction='mean'):
+  """Computes the AUC group by the key separately.
+
+  Args:
+    labels: A `Tensor` whose shape matches `predictions`. Will be cast to
+      `bool`.
+    predictions: A floating point `Tensor` of arbitrary shape and whose values
+      are in the range `[0, 1]`.
+    keys: keys to be group by, A int or string `Tensor` whose shape matches `predictions`.
+    reduction: reduction metric for auc of different keys
+      * "mean": simple mean of different keys
+      * "mean_by_sample_num": weighted mean with sample num of different keys
+      * "mean_by_positive_num": weighted mean with positive sample num of different keys
+  """
+  assert reduction in ['mean', 'mean_by_sample_num', 'mean_by_positive_num'], \
+      'reduction method must in mean | mean_by_sample_num | mean_by_positive_num'
+  separated_label = defaultdict(list)
+  separated_prediction = defaultdict(list)
+  separated_weights = defaultdict(int)
+
+  def update_pyfunc(labels, predictions, keys):
+    for label, prediction, key in zip(labels, predictions, keys):
+      separated_label[key].append(label)
+      separated_prediction[key].append(prediction)
+      if reduction == 'mean':
+        separated_weights[key] = 1
+      elif reduction == 'mean_by_sample_num':
+        separated_weights[key] += 1
+      elif reduction == 'mean_by_positive_num':
+        separated_weights[key] += label
+
+  def value_pyfunc():
+    metrics = []
+    weights = []
+    for key in separated_label.keys():
+      per_label = np.asarray(separated_label[key]).reshape([-1])
+      per_prediction = np.asarray(separated_prediction[key]).reshape([-1])
+      if np.all(per_label == 1) or np.all(per_label == 0):
+        continue
+      metric = sklearn_metrics.roc_auc_score(per_label, per_prediction)
+      metrics.append(metric)
+      weights.append(separated_weights[key])
+    if len(metrics) > 0:
+      return np.average(metrics, weights=weights).astype(np.float32)
+    else:
+      return np.float32(0.0)
+
+  update_op = tf.py_func(update_pyfunc, [labels, predictions, keys], [])
+  value_op = tf.py_func(value_pyfunc, [], tf.float32)
+  return value_op, update_op
+
+
+def fast_auc(labels, predictions, name, num_thresholds=1e5):
+  num_thresholds = int(num_thresholds)
+
+  def value_pyfunc(pos_neg_arr, total_pos_neg):
+    partial_sum_pos = 0
+    auc = 0
+    total_neg = total_pos_neg[0]
+    total_pos = total_pos_neg[1]
+    for i in range(num_thresholds + 1):
+      partial_sum_pos += pos_neg_arr[1][i]
+      auc += (total_pos - partial_sum_pos) * pos_neg_arr[0][i] * 2
+      auc += pos_neg_arr[0][i] * pos_neg_arr[1][i]
+    auc = np.double(auc) / np.double(total_pos * total_neg * 2)
+    logging.info('fast_auc[%s]: total_pos=%d total_neg=%d total=%d' %
+                 (name, total_pos, total_neg, total_pos + total_neg))
+    return np.float32(auc)
+
+  with variable_scope.variable_scope(name_or_scope=name), tf.name_scope(name):
+    neg_pos_var = variable_scope.get_variable(
+        name='neg_pos_cnt',
+        shape=[2, num_thresholds + 1],
+        trainable=False,
+        collections=[tf.GraphKeys.METRIC_VARIABLES],
+        initializer=tf.zeros_initializer(),
+        dtype=tf.int64)
+    total_var = variable_scope.get_variable(
+        name='total_cnt',
+        shape=[2],
+        trainable=False,
+        collections=[tf.GraphKeys.METRIC_VARIABLES],
+        initializer=tf.zeros_initializer(),
+        dtype=tf.int64)
+    pred_bins = math_ops.cast(predictions * num_thresholds, dtype=tf.int32)
+    labels = math_ops.cast(labels, dtype=tf.int32)
+    labels = array_ops.reshape(labels, [-1, 1])
+    pred_bins = array_ops.reshape(pred_bins, [-1, 1])
+    update_op0 = state_ops.scatter_nd_add(
+        neg_pos_var, tf.concat([labels, pred_bins], axis=1),
+        array_ops.ones(tf.shape(labels)[0], dtype=tf.int64))
+    total_pos = math_ops.reduce_sum(labels)
+    total_neg = array_ops.shape(labels)[0] - total_pos
+    total_add = math_ops.cast(tf.stack([total_neg, total_pos]), dtype=tf.int64)
+    update_op1 = state_ops.assign_add(total_var, total_add)
+    return tf.py_func(value_pyfunc, [neg_pos_var, total_var],
+                      tf.float32), tf.group([update_op0, update_op1])
+
+
+def _distribute_separated_auc_impl(labels,
+                                   predictions,
+                                   keys,
+                                   reduction='mean',
+                                   metric_name='sepatated_auc'):
+  """Computes the AUC group by the key separately.
+
+  Args:
+    labels: A `Tensor` whose shape matches `predictions`. Will be cast to
+      `bool`.
+    predictions: A floating point `Tensor` of arbitrary shape and whose values
+      are in the range `[0, 1]`.
+    keys: keys to be group by, A int or string `Tensor` whose shape matches `predictions`.
+    reduction: reduction metric for auc of different keys
+    metric_name: the name of compute metric
+      * "mean": simple mean of different keys
+      * "mean_by_sample_num": weighted mean with sample num of different keys
+      * "mean_by_positive_num": weighted mean with positive sample num of different keys
+  """
+  assert reduction in ['mean', 'mean_by_sample_num', 'mean_by_positive_num'], \
+      'reduction method must in mean | mean_by_sample_num | mean_by_positive_num'
+  separated_label = defaultdict(list)
+  separated_prediction = defaultdict(list)
+  separated_weights = defaultdict(int)
+  tf_config = json.loads(os.environ['TF_CONFIG'])
+  cur_job_name = tf_config['task']['type']
+  cur_task_index, task_num = get_task_index_and_num()
+  cur_work_device = 'job_' + cur_job_name + '__' + 'task_' + str(cur_task_index)
+  eval_tmp_results_dir = os.environ['eval_tmp_results_dir']
+  assert tf.gfile.IsDirectory(
+      eval_tmp_results_dir), 'eval_tmp_results_dir not exists'
+
+  def update_pyfunc(labels, predictions, keys):
+    for label, prediction, key in zip(labels, predictions, keys):
+      key = str(key)
+      separated_label[key].append(label.item())
+      separated_prediction[key].append(prediction.item())
+      if reduction == 'mean':
+        separated_weights[key] = 1
+      elif reduction == 'mean_by_sample_num':
+        separated_weights[key] += 1
+      elif reduction == 'mean_by_positive_num':
+        separated_weights[key] += label.item()
+    for name, data in zip(
+        ['separated_label', 'separated_prediction', 'separated_weights'],
+        [separated_label, separated_prediction, separated_weights]):
+      cur_json_name = metric_name + '__' + cur_work_device + '__' + name + '.json'
+      cur_json_path = os.path.join(eval_tmp_results_dir, cur_json_name)
+      save_data_to_json_path(cur_json_path, data)
+
+  def value_pyfunc():
+    for task_i in range(1, task_num):
+      work_device_i = 'job_worker__task_' + str(task_i)
+      for name in [
+          'separated_label', 'separated_prediction', 'separated_weights'
+      ]:
+        json_name_i = metric_name + '__' + work_device_i + '__' + name + '.json'
+        json_path_i = os.path.join(eval_tmp_results_dir, json_name_i)
+        data_i = read_data_from_json_path(json_path_i)
+        if (name == 'separated_label'):
+          separated_label.update({
+              key: separated_label.get(key, []) + data_i.get(key, [])
+              for key in set(
+                  list(separated_label.keys()) + list(data_i.keys()))
+          })
+        elif (name == 'separated_prediction'):
+          separated_prediction.update({
+              key: separated_prediction.get(key, []) + data_i.get(key, [])
+              for key in set(
+                  list(separated_prediction.keys()) + list(data_i.keys()))
+          })
+        elif (name == 'separated_weights'):
+          if reduction == 'mean':
+            separated_weights.update(data_i)
+          else:
+            separated_weights.update({
+                key: separated_weights.get(key, 0) + data_i.get(key, 0)
+                for key in set(
+                    list(separated_weights.keys()) + list(data_i.keys()))
+            })
+        else:
+          assert False, 'Not supported name {}'.format(name)
+    metrics = []
+    weights = []
+    for key in separated_label.keys():
+      per_label = np.asarray(separated_label[key]).reshape([-1])
+      per_prediction = np.asarray(separated_prediction[key]).reshape([-1])
+      if np.all(per_label == 1) or np.all(per_label == 0):
+        continue
+      metric = sklearn_metrics.roc_auc_score(per_label, per_prediction)
+      metrics.append(metric)
+      weights.append(separated_weights[key])
+    if len(metrics) > 0:
+      return np.average(metrics, weights=weights).astype(np.float32)
+    else:
+      return np.float32(0.0)
+
+  update_op = tf.py_func(update_pyfunc, [labels, predictions, keys], [])
+  value_op = tf.py_func(value_pyfunc, [], tf.float32)
+  return value_op, update_op
+
+
+def gauc(labels, predictions, uids, reduction='mean'):
+  """Computes the AUC group by user separately.
+
+  Args:
+    labels: A `Tensor` whose shape matches `predictions`. Will be cast to
+      `bool`.
+    predictions: A floating point `Tensor` of arbitrary shape and whose values
+      are in the range `[0, 1]`.
+    uids: user ids, A int or string `Tensor` whose shape matches `predictions`.
+    reduction: reduction method for auc of different users
+      * "mean": simple mean of different users
+      * "mean_by_sample_num": weighted mean with sample num of different users
+      * "mean_by_positive_num": weighted mean with positive sample num of different users
+  """
+  if os.environ.get('distribute_eval') == 'True':
+    return _distribute_separated_auc_impl(
+        labels, predictions, uids, reduction, metric_name='gauc')
+  return _separated_auc_impl(labels, predictions, uids, reduction)
+
+
+def session_auc(labels, predictions, session_ids, reduction='mean'):
+  """Computes the AUC group by session separately.
+
+  Args:
+    labels: A `Tensor` whose shape matches `predictions`. Will be cast to
+      `bool`.
+    predictions: A floating point `Tensor` of arbitrary shape and whose values
+      are in the range `[0, 1]`.
+    session_ids: session ids, A int or string `Tensor` whose shape matches `predictions`.
+    reduction: reduction method for auc of different sessions
+      * "mean": simple mean of different sessions
+      * "mean_by_sample_num": weighted mean with sample num of different sessions
+      * "mean_by_positive_num": weighted mean with positive sample num of different sessions
+  """
+  if os.environ.get('distribute_eval') == 'True':
+    return _distribute_separated_auc_impl(
+        labels, predictions, session_ids, reduction, metric_name='session_auc')
+  return _separated_auc_impl(labels, predictions, session_ids, reduction)
+
+
+def metric_learning_recall_at_k(k,
+                                embeddings,
+                                labels,
+                                session_ids=None,
+                                embed_normed=False):
+  """Computes the recall_at_k metric for metric learning.
+
+  Args:
+    k: a scalar of int, or a tuple of ints
+    embeddings: the output of last hidden layer, a tf.float32 `Tensor` with shape [batch_size, embedding_size]
+    labels: a `Tensor` with shape [batch_size]
+    session_ids: session ids, a `Tensor` with shape [batch_size]
+    embed_normed: indicator of whether the input embeddings are l2_normalized
+  """
+  from easy_rec.python.core.easyrec_metrics import metrics_tf
+  # make sure embedding should be l2-normalized
+  if not embed_normed:
+    embeddings = tf.nn.l2_normalize(embeddings, axis=1)
+  embed_shape = get_shape_list(embeddings)
+  batch_size = embed_shape[0]
+  sim_mat = tf.matmul(embeddings, embeddings, transpose_b=True)
+  sim_mat = sim_mat - tf.eye(batch_size) * 2.0
+  indices_not_equal = tf.logical_not(tf.eye(batch_size, dtype=tf.bool))
+  # Uses broadcasting where the 1st argument has shape (1, batch_size) and the 2nd (batch_size, 1)
+  labels_equal = tf.equal(tf.expand_dims(labels, 0), tf.expand_dims(labels, 1))
+  if session_ids is not None and session_ids is not labels:
+    sessions_equal = tf.equal(
+        tf.expand_dims(session_ids, 0), tf.expand_dims(session_ids, 1))
+    labels_equal = tf.logical_and(sessions_equal, labels_equal)
+  mask = tf.logical_and(indices_not_equal, labels_equal)
+  mask_pos = tf.where(
+      mask, sim_mat,
+      -array_ops.ones_like(sim_mat))  # shape: (batch_size, batch_size)
+  if isinstance(k, int):
+    _, pos_top_k_idx = tf.nn.top_k(mask_pos, k)  # shape: (batch_size, k)
+    return metrics_tf.recall_at_k(
+        labels=tf.to_int64(pos_top_k_idx), predictions=sim_mat, k=k)
+  if any((isinstance(k, list), isinstance(k, tuple), isinstance(k, set))):
+    metrics = {}
+    for kk in k:
+      if kk < 1:
+        continue
+      _, pos_top_k_idx = tf.nn.top_k(mask_pos, kk)
+      metrics['recall@' + str(kk)] = metrics_tf.recall_at_k(
+          labels=tf.to_int64(pos_top_k_idx), predictions=sim_mat, k=kk)
+    return metrics
+  else:
+    raise ValueError('k should be a `int` or a list/tuple/set of int.')
+
+
+def metric_learning_average_precision_at_k(k,
+                                           embeddings,
+                                           labels,
+                                           session_ids=None,
+                                           embed_normed=False):
+  from easy_rec.python.core.easyrec_metrics import metrics_tf
+  # make sure embedding should be l2-normalized
+  if not embed_normed:
+    embeddings = tf.nn.l2_normalize(embeddings, axis=1)
+  embed_shape = get_shape_list(embeddings)
+  batch_size = embed_shape[0]
+  sim_mat = tf.matmul(embeddings, embeddings, transpose_b=True)
+  sim_mat = sim_mat - tf.eye(batch_size) * 2.0
+  mask = tf.equal(tf.expand_dims(labels, 0), tf.expand_dims(labels, 1))
+  if session_ids is not None and session_ids is not labels:
+    sessions_equal = tf.equal(
+        tf.expand_dims(session_ids, 0), tf.expand_dims(session_ids, 1))
+    mask = tf.logical_and(sessions_equal, mask)
+  label_indices = _get_matrix_mask_indices(mask)
+  if isinstance(k, int):
+    return metrics_tf.average_precision_at_k(label_indices, sim_mat, k)
+  if any((isinstance(k, list), isinstance(k, tuple), isinstance(k, set))):
+    metrics = {}
+    for kk in k:
+      if kk < 1:
+        continue
+      metrics['MAP@' + str(kk)] = metrics_tf.average_precision_at_k(
+          label_indices, sim_mat, kk)
+    return metrics
+  else:
+    raise ValueError('k should be a `int` or a list/tuple/set of int.')
+
+
+def _get_matrix_mask_indices(matrix, num_rows=None):
+  if num_rows is None:
+    num_rows = get_shape_list(matrix)[0]
+  indices = tf.where(matrix)
+  num_indices = tf.shape(indices)[0]
+  elem_per_row = tf.bincount(
+      tf.cast(indices[:, 0], tf.int32), minlength=num_rows)
+  max_elem_per_row = tf.reduce_max(elem_per_row)
+  row_start = tf.concat([[0], tf.cumsum(elem_per_row[:-1])], axis=0)
+  r = tf.range(max_elem_per_row)
+  idx = tf.expand_dims(row_start, 1) + r
+  idx = tf.minimum(idx, num_indices - 1)
+  result = tf.gather(indices[:, 1], idx)
+  # replace invalid elements with -1
+  result = tf.where(
+      tf.expand_dims(elem_per_row, 1) > r, result, -array_ops.ones_like(result))
+  max_index_per_row = tf.reduce_max(result, axis=1, keepdims=True)
+  max_index_per_row = tf.tile(max_index_per_row, [1, max_elem_per_row])
+  result = tf.where(result >= 0, result, max_index_per_row)
+  return result