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

easy_rec/python/model/rank_model.py

@@ -0,0 +1,485 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+from tensorflow.python.ops import math_ops
+
+from easy_rec.python.builders import loss_builder
+from easy_rec.python.model.easy_rec_model import EasyRecModel
+from easy_rec.python.protos.loss_pb2 import LossType
+
+from easy_rec.python.loss.zero_inflated_lognormal import zero_inflated_lognormal_pred  # NOQA
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class RankModel(EasyRecModel):
+
+  def __init__(self,
+               model_config,
+               feature_configs,
+               features,
+               labels=None,
+               is_training=False):
+    super(RankModel, self).__init__(model_config, feature_configs, features,
+                                    labels, is_training)
+    self._loss_type = self._model_config.loss_type
+    self._num_class = self._model_config.num_class
+    self._losses = self._model_config.losses
+    if self._labels is not None:
+      if model_config.HasField('label_name'):
+        self._label_name = model_config.label_name
+      else:
+        self._label_name = list(self._labels.keys())[0]
+    self._outputs = []
+
+  def build_predict_graph(self):
+    if not self.has_backbone:
+      raise NotImplementedError(
+          'method `build_predict_graph` must be implemented when backbone network do not exits'
+      )
+    model = self._model_config.WhichOneof('model')
+    assert model == 'model_params', '`model_params` must be configured'
+    config = self._model_config.model_params
+    for out in config.outputs:
+      self._outputs.append(out)
+
+    output = self.backbone
+    if int(output.shape[-1]) != self._num_class:
+      logging.info('add head logits layer for rank model')
+      output = tf.layers.dense(output, self._num_class, name='output')
+    self._add_to_prediction_dict(output)
+    return self._prediction_dict
+
+  def _output_to_prediction_impl(self,
+                                 output,
+                                 loss_type,
+                                 num_class=1,
+                                 suffix=''):
+    prediction_dict = {}
+    binary_loss_type = {
+        LossType.F1_REWEIGHTED_LOSS, LossType.PAIR_WISE_LOSS,
+        LossType.BINARY_FOCAL_LOSS, LossType.PAIRWISE_FOCAL_LOSS,
+        LossType.LISTWISE_RANK_LOSS, LossType.PAIRWISE_HINGE_LOSS,
+        LossType.PAIRWISE_LOGISTIC_LOSS, LossType.BINARY_CROSS_ENTROPY_LOSS,
+        LossType.LISTWISE_DISTILL_LOSS
+    }
+    if loss_type in binary_loss_type:
+      assert num_class == 1, 'num_class must be 1 when loss type is %s' % loss_type.name
+      output = tf.squeeze(output, axis=1)
+      probs = tf.sigmoid(output)
+      tf.summary.scalar('prediction/probs', tf.reduce_mean(probs))
+      prediction_dict['logits' + suffix] = output
+      prediction_dict['probs' + suffix] = probs
+    elif loss_type == LossType.JRC_LOSS:
+      assert num_class == 2, 'num_class must be 2 when loss type is JRC_LOSS'
+      probs = tf.nn.softmax(output, axis=1)
+      tf.summary.scalar('prediction/probs', tf.reduce_mean(probs[:, 1]))
+      prediction_dict['logits' + suffix] = output
+      prediction_dict['pos_logits' + suffix] = output[:, 1]
+      prediction_dict['probs' + suffix] = probs[:, 1]
+    elif loss_type == LossType.ZILN_LOSS:
+      assert num_class == 3, 'num_class must be 3 when loss type is ZILN_LOSS'
+      probs, preds = zero_inflated_lognormal_pred(output)
+      tf.summary.scalar('prediction/probs', tf.reduce_mean(probs))
+      tf.summary.scalar('prediction/y', tf.reduce_mean(preds))
+      prediction_dict['logits' + suffix] = output
+      prediction_dict['probs' + suffix] = probs
+      prediction_dict['y' + suffix] = preds
+    elif loss_type == LossType.CLASSIFICATION:
+      if num_class == 1:
+        output = tf.squeeze(output, axis=1)
+        probs = tf.sigmoid(output)
+        tf.summary.scalar('prediction/probs', tf.reduce_mean(probs))
+        prediction_dict['logits' + suffix] = output
+        prediction_dict['probs' + suffix] = probs
+      else:
+        probs = tf.nn.softmax(output, axis=1)
+        prediction_dict['logits' + suffix] = output
+        prediction_dict['logits' + suffix + '_1'] = output[:, 1]
+        prediction_dict['probs' + suffix] = probs
+        prediction_dict['probs' + suffix + '_1'] = probs[:, 1]
+        prediction_dict['logits' + suffix + '_y'] = math_ops.reduce_max(
+            output, axis=1)
+        prediction_dict['probs' + suffix + '_y'] = math_ops.reduce_max(
+            probs, axis=1)
+        prediction_dict['y' + suffix] = tf.argmax(output, axis=1)
+    elif loss_type == LossType.L2_LOSS:
+      output = tf.squeeze(output, axis=1)
+      prediction_dict['y' + suffix] = output
+    elif loss_type == LossType.SIGMOID_L2_LOSS:
+      output = tf.squeeze(output, axis=1)
+      prediction_dict['y' + suffix] = tf.sigmoid(output)
+    return prediction_dict
+
+  def _add_to_prediction_dict(self, output):
+    if len(self._losses) == 0:
+      prediction_dict = self._output_to_prediction_impl(
+          output, loss_type=self._loss_type, num_class=self._num_class)
+      self._prediction_dict.update(prediction_dict)
+    else:
+      for loss in self._losses:
+        prediction_dict = self._output_to_prediction_impl(
+            output, loss_type=loss.loss_type, num_class=self._num_class)
+        self._prediction_dict.update(prediction_dict)
+
+  def build_rtp_output_dict(self):
+    """Forward tensor as `rank_predict`, which is a special node for RTP."""
+    outputs = {}
+    outputs.update(super(RankModel, self).build_rtp_output_dict())
+    rank_predict = None
+    try:
+      op = tf.get_default_graph().get_operation_by_name('rank_predict')
+      if len(op.outputs) != 1:
+        raise ValueError(
+            ('failed to build RTP rank_predict output: op {}[{}] has output ' +
+             'size {}, however 1 is expected.').format(op.name, op.type,
+                                                       len(op.outputs)))
+      rank_predict = op.outputs[0]
+    except KeyError:
+      forwarded = None
+      loss_types = {self._loss_type}
+      if len(self._losses) > 0:
+        loss_types = {loss.loss_type for loss in self._losses}
+      binary_loss_set = {
+          LossType.CLASSIFICATION, LossType.F1_REWEIGHTED_LOSS,
+          LossType.PAIR_WISE_LOSS, LossType.BINARY_FOCAL_LOSS,
+          LossType.PAIRWISE_FOCAL_LOSS, LossType.PAIRWISE_LOGISTIC_LOSS,
+          LossType.JRC_LOSS, LossType.LISTWISE_DISTILL_LOSS,
+          LossType.LISTWISE_RANK_LOSS
+      }
+      if loss_types & binary_loss_set:
+        if 'probs' in self._prediction_dict:
+          forwarded = self._prediction_dict['probs']
+        else:
+          raise ValueError(
+              'failed to build RTP rank_predict output: classification model ' +
+              "expect 'probs' prediction, which is not found. Please check if" +
+              ' build_predict_graph() is called.')
+      elif loss_types & {
+          LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS, LossType.ZILN_LOSS
+      }:
+        if 'y' in self._prediction_dict:
+          forwarded = self._prediction_dict['y']
+        else:
+          raise ValueError(
+              'failed to build RTP rank_predict output: regression model expect'
+              +
+              "'y' prediction, which is not found. Please check if build_predic"
+              + 't_graph() is called.')
+      else:
+        logging.warning(
+            'failed to build RTP rank_predict: unsupported loss type {}'.format(
+                loss_types))
+      if forwarded is not None:
+        rank_predict = tf.identity(forwarded, name='rank_predict')
+    if rank_predict is not None:
+      outputs['rank_predict'] = rank_predict
+    return outputs
+
+  def _build_loss_impl(self,
+                       loss_type,
+                       label_name,
+                       loss_weight=1.0,
+                       num_class=1,
+                       suffix='',
+                       loss_name='',
+                       loss_param=None):
+    loss_dict = {}
+    binary_loss_type = {
+        LossType.F1_REWEIGHTED_LOSS, LossType.PAIR_WISE_LOSS,
+        LossType.BINARY_FOCAL_LOSS, LossType.PAIRWISE_FOCAL_LOSS,
+        LossType.LISTWISE_RANK_LOSS, LossType.PAIRWISE_HINGE_LOSS,
+        LossType.PAIRWISE_LOGISTIC_LOSS, LossType.JRC_LOSS,
+        LossType.LISTWISE_DISTILL_LOSS, LossType.ZILN_LOSS
+    }
+    if loss_type in {
+        LossType.CLASSIFICATION, LossType.BINARY_CROSS_ENTROPY_LOSS
+    }:
+      loss_name = loss_name if loss_name else 'cross_entropy_loss' + suffix
+      pred = self._prediction_dict['logits' + suffix]
+    elif loss_type in binary_loss_type:
+      if not loss_name:
+        loss_name = LossType.Name(loss_type).lower() + suffix
+      else:
+        loss_name = loss_name + suffix
+      pred = self._prediction_dict['logits' + suffix]
+    elif loss_type in [LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS]:
+      loss_name = loss_name if loss_name else 'l2_loss' + suffix
+      pred = self._prediction_dict['y' + suffix]
+    else:
+      raise ValueError('invalid loss type: %s' % LossType.Name(loss_type))
+
+    tf.summary.scalar('labels/%s' % label_name,
+                      tf.reduce_mean(tf.to_float(self._labels[label_name])))
+    kwargs = {'loss_name': loss_name}
+    if loss_param is not None:
+      if hasattr(loss_param, 'session_name'):
+        kwargs['session_ids'] = self._feature_dict[loss_param.session_name]
+    loss_dict[loss_name] = loss_builder.build(
+        loss_type,
+        self._labels[label_name],
+        pred,
+        loss_weight,
+        num_class,
+        loss_param=loss_param,
+        **kwargs)
+    return loss_dict
+
+  def build_loss_graph(self):
+    loss_dict = {}
+    with tf.name_scope('loss'):
+      if len(self._losses) == 0:
+        loss_dict = self._build_loss_impl(
+            self._loss_type,
+            label_name=self._label_name,
+            loss_weight=self._sample_weight,
+            num_class=self._num_class)
+      else:
+        strategy = self._base_model_config.loss_weight_strategy
+        loss_weight = [1.0]
+        if strategy == self._base_model_config.Random and len(self._losses) > 1:
+          weights = tf.random_normal([len(self._losses)])
+          loss_weight = tf.nn.softmax(weights)
+        for i, loss in enumerate(self._losses):
+          loss_param = loss.WhichOneof('loss_param')
+          if loss_param is not None:
+            loss_param = getattr(loss, loss_param)
+          loss_ops = self._build_loss_impl(
+              loss.loss_type,
+              label_name=self._label_name,
+              loss_weight=self._sample_weight,
+              num_class=self._num_class,
+              loss_name=loss.loss_name,
+              loss_param=loss_param)
+          for loss_name, loss_value in loss_ops.items():
+            if strategy == self._base_model_config.Fixed:
+              loss_dict[loss_name] = loss_value * loss.weight
+            elif strategy == self._base_model_config.Uncertainty:
+              if loss.learn_loss_weight:
+                uncertainty = tf.Variable(
+                    0, name='%s_loss_weight' % loss_name, dtype=tf.float32)
+                tf.summary.scalar('%s_uncertainty' % loss_name, uncertainty)
+                if loss.loss_type in {
+                    LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS
+                }:
+                  loss_dict[loss_name] = 0.5 * tf.exp(
+                      -uncertainty) * loss_value + 0.5 * uncertainty
+                else:
+                  loss_dict[loss_name] = tf.exp(
+                      -uncertainty) * loss_value + 0.5 * uncertainty
+              else:
+                loss_dict[loss_name] = loss_value * loss.weight
+            elif strategy == self._base_model_config.Random:
+              loss_dict[loss_name] = loss_value * loss_weight[i]
+            else:
+              raise ValueError('Unsupported loss weight strategy: ' +
+                               strategy.Name)
+      self._loss_dict.update(loss_dict)
+      # build kd loss
+      kd_loss_dict = loss_builder.build_kd_loss(self.kd, self._prediction_dict,
+                                                self._labels,
+                                                self._feature_dict)
+      self._loss_dict.update(kd_loss_dict)
+    return self._loss_dict
+
+  def _build_metric_impl(self,
+                         metric,
+                         loss_type,
+                         label_name,
+                         num_class=1,
+                         suffix=''):
+    if not isinstance(loss_type, set):
+      loss_type = {loss_type}
+    from easy_rec.python.core.easyrec_metrics import metrics_tf
+    from easy_rec.python.core import metrics as metrics_lib
+    binary_loss_set = {
+        LossType.CLASSIFICATION, LossType.F1_REWEIGHTED_LOSS,
+        LossType.PAIR_WISE_LOSS, LossType.BINARY_FOCAL_LOSS,
+        LossType.PAIRWISE_FOCAL_LOSS, LossType.PAIRWISE_LOGISTIC_LOSS,
+        LossType.JRC_LOSS, LossType.LISTWISE_DISTILL_LOSS,
+        LossType.LISTWISE_RANK_LOSS, LossType.ZILN_LOSS
+    }
+    metric_dict = {}
+    if metric.WhichOneof('metric') == 'auc':
+      assert loss_type & binary_loss_set
+      if num_class == 1 or loss_type & {LossType.JRC_LOSS, LossType.ZILN_LOSS}:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['auc' + suffix] = metrics_tf.auc(
+            label,
+            self._prediction_dict['probs' + suffix],
+            num_thresholds=metric.auc.num_thresholds)
+      elif num_class == 2:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['auc' + suffix] = metrics_tf.auc(
+            label,
+            self._prediction_dict['probs' + suffix][:, 1],
+            num_thresholds=metric.auc.num_thresholds)
+      else:
+        raise ValueError('Wrong class number')
+    elif metric.WhichOneof('metric') == 'gauc':
+      assert loss_type & binary_loss_set
+      if num_class == 1 or loss_type & {LossType.JRC_LOSS, LossType.ZILN_LOSS}:
+        label = tf.to_int64(self._labels[label_name])
+        uids = self._feature_dict[metric.gauc.uid_field]
+        if isinstance(uids, tf.sparse.SparseTensor):
+          uids = tf.sparse_to_dense(
+              uids.indices, uids.dense_shape, uids.values, default_value='')
+          uids = tf.reshape(uids, [-1])
+        metric_dict['gauc' + suffix] = metrics_lib.gauc(
+            label,
+            self._prediction_dict['probs' + suffix],
+            uids=uids,
+            reduction=metric.gauc.reduction)
+      elif num_class == 2:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['gauc' + suffix] = metrics_lib.gauc(
+            label,
+            self._prediction_dict['probs' + suffix][:, 1],
+            uids=self._feature_dict[metric.gauc.uid_field],
+            reduction=metric.gauc.reduction)
+      else:
+        raise ValueError('Wrong class number')
+    elif metric.WhichOneof('metric') == 'session_auc':
+      assert loss_type & binary_loss_set
+      if num_class == 1 or loss_type & {LossType.JRC_LOSS, LossType.ZILN_LOSS}:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['session_auc' + suffix] = metrics_lib.session_auc(
+            label,
+            self._prediction_dict['probs' + suffix],
+            session_ids=self._feature_dict[metric.session_auc.session_id_field],
+            reduction=metric.session_auc.reduction)
+      elif num_class == 2:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['session_auc' + suffix] = metrics_lib.session_auc(
+            label,
+            self._prediction_dict['probs' + suffix][:, 1],
+            session_ids=self._feature_dict[metric.session_auc.session_id_field],
+            reduction=metric.session_auc.reduction)
+      else:
+        raise ValueError('Wrong class number')
+    elif metric.WhichOneof('metric') == 'max_f1':
+      assert loss_type & binary_loss_set
+      if num_class == 1 or loss_type & {LossType.JRC_LOSS, LossType.ZILN_LOSS}:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['max_f1' + suffix] = metrics_lib.max_f1(
+            label, self._prediction_dict['logits' + suffix])
+      elif num_class == 2:
+        label = tf.to_int64(self._labels[label_name])
+        metric_dict['max_f1' + suffix] = metrics_lib.max_f1(
+            label, self._prediction_dict['logits' + suffix][:, 1])
+      else:
+        raise ValueError('Wrong class number')
+    elif metric.WhichOneof('metric') == 'recall_at_topk':
+      assert loss_type & binary_loss_set
+      assert num_class > 1
+      label = tf.to_int64(self._labels[label_name])
+      metric_dict['recall_at_topk' + suffix] = metrics_tf.recall_at_k(
+          label, self._prediction_dict['logits' + suffix],
+          metric.recall_at_topk.topk)
+    elif metric.WhichOneof('metric') == 'mean_absolute_error':
+      label = tf.to_float(self._labels[label_name])
+      if loss_type & {
+          LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS, LossType.ZILN_LOSS
+      }:
+        metric_dict['mean_absolute_error' +
+                    suffix] = metrics_tf.mean_absolute_error(
+                        label, self._prediction_dict['y' + suffix])
+      elif loss_type & {LossType.CLASSIFICATION} and num_class == 1:
+        metric_dict['mean_absolute_error' +
+                    suffix] = metrics_tf.mean_absolute_error(
+                        label, self._prediction_dict['probs' + suffix])
+      else:
+        assert False, 'mean_absolute_error is not supported for this model'
+    elif metric.WhichOneof('metric') == 'mean_squared_error':
+      label = tf.to_float(self._labels[label_name])
+      if loss_type & {
+          LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS, LossType.ZILN_LOSS
+      }:
+        metric_dict['mean_squared_error' +
+                    suffix] = metrics_tf.mean_squared_error(
+                        label, self._prediction_dict['y' + suffix])
+      elif num_class == 1 and loss_type & binary_loss_set:
+        metric_dict['mean_squared_error' +
+                    suffix] = metrics_tf.mean_squared_error(
+                        label, self._prediction_dict['probs' + suffix])
+      else:
+        assert False, 'mean_squared_error is not supported for this model'
+    elif metric.WhichOneof('metric') == 'root_mean_squared_error':
+      label = tf.to_float(self._labels[label_name])
+      if loss_type & {
+          LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS, LossType.ZILN_LOSS
+      }:
+        metric_dict['root_mean_squared_error' +
+                    suffix] = metrics_tf.root_mean_squared_error(
+                        label, self._prediction_dict['y' + suffix])
+      elif loss_type & {LossType.CLASSIFICATION} and num_class == 1:
+        metric_dict['root_mean_squared_error' +
+                    suffix] = metrics_tf.root_mean_squared_error(
+                        label, self._prediction_dict['probs' + suffix])
+      else:
+        assert False, 'root_mean_squared_error is not supported for this model'
+    elif metric.WhichOneof('metric') == 'accuracy':
+      assert loss_type & {LossType.CLASSIFICATION}
+      assert num_class > 1
+      label = tf.to_int64(self._labels[label_name])
+      metric_dict['accuracy' + suffix] = metrics_tf.accuracy(
+          label, self._prediction_dict['y' + suffix])
+    return metric_dict
+
+  def build_metric_graph(self, eval_config):
+    loss_types = {self._loss_type}
+    if len(self._losses) > 0:
+      loss_types = {loss.loss_type for loss in self._losses}
+    for metric in eval_config.metrics_set:
+      self._metric_dict.update(
+          self._build_metric_impl(
+              metric,
+              loss_type=loss_types,
+              label_name=self._label_name,
+              num_class=self._num_class))
+    return self._metric_dict
+
+  def _get_outputs_impl(self, loss_type, num_class=1, suffix=''):
+    binary_loss_set = {
+        LossType.F1_REWEIGHTED_LOSS, LossType.PAIR_WISE_LOSS,
+        LossType.BINARY_FOCAL_LOSS, LossType.PAIRWISE_FOCAL_LOSS,
+        LossType.LISTWISE_RANK_LOSS, LossType.PAIRWISE_HINGE_LOSS,
+        LossType.PAIRWISE_LOGISTIC_LOSS, LossType.LISTWISE_DISTILL_LOSS
+    }
+    if loss_type in binary_loss_set:
+      return ['probs' + suffix, 'logits' + suffix]
+    if loss_type == LossType.JRC_LOSS:
+      return ['probs' + suffix, 'pos_logits' + suffix]
+    if loss_type == LossType.ZILN_LOSS:
+      return ['probs' + suffix, 'y' + suffix, 'logits' + suffix]
+    if loss_type == LossType.CLASSIFICATION:
+      if num_class == 1:
+        return ['probs' + suffix, 'logits' + suffix]
+      else:
+        return [
+            'y' + suffix, 'probs' + suffix, 'logits' + suffix,
+            'probs' + suffix + '_y', 'logits' + suffix + '_y',
+            'probs' + suffix + '_1', 'logits' + suffix + '_1'
+        ]
+    elif loss_type in [LossType.L2_LOSS, LossType.SIGMOID_L2_LOSS]:
+      return ['y' + suffix]
+    else:
+      raise ValueError('invalid loss type: %s' % LossType.Name(loss_type))
+
+  def get_outputs(self):
+    if len(self._losses) == 0:
+      outputs = self._get_outputs_impl(self._loss_type, self._num_class)
+      if self._outputs:
+        outputs.extend(self._outputs)
+      return list(set(outputs))
+
+    all_outputs = []
+    if self._outputs:
+      all_outputs.extend(self._outputs)
+    for loss in self._losses:
+      outputs = self._get_outputs_impl(loss.loss_type, self._num_class)
+      all_outputs.extend(outputs)
+    return list(set(all_outputs))

easy_rec/python/model/rocket_launching.py

@@ -0,0 +1,203 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+
+from easy_rec.python.builders import loss_builder
+from easy_rec.python.layers import dnn
+from easy_rec.python.model.rank_model import RankModel
+from easy_rec.python.protos.loss_pb2 import LossType
+from easy_rec.python.protos.simi_pb2 import Similarity
+
+from easy_rec.python.protos.rocket_launching_pb2 import RocketLaunching as RocketLaunchingConfig  # NOQA
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class RocketLaunching(RankModel):
+
+  def __init__(self,
+               model_config,
+               feature_configs,
+               features,
+               labels=None,
+               is_training=False):
+    super(RocketLaunching, self).__init__(model_config, feature_configs,
+                                          features, labels, is_training)
+    assert self._model_config.WhichOneof('model') == 'rocket_launching', \
+        'invalid model config: %s' % self._model_config.WhichOneof('model')
+    self._model_config = self._model_config.rocket_launching
+    assert isinstance(self._model_config, RocketLaunchingConfig)
+    if self._labels is not None:
+      self._label_name = list(self._labels.keys())[0]
+
+    self._features, _ = self._input_layer(self._feature_dict, 'all')
+
+  def sim(self, feature_emb1, feature_emb2):
+    emb1_emb2_sim = tf.reduce_sum(
+        tf.multiply(feature_emb1, feature_emb2), axis=1, keepdims=True)
+    return emb1_emb2_sim
+
+  def norm(self, fea):
+    fea_norm = tf.nn.l2_normalize(fea, axis=1)
+    return fea_norm
+
+  def feature_based_sim(self, feature_based_distillation, i, j):
+    booster_feature_no_gradient = tf.stop_gradient(
+        self.booster_feature['hidden_layer' + str(j)])
+    if feature_based_distillation == Similarity.COSINE:
+      booster_feature_no_gradient_norm = self.norm(booster_feature_no_gradient)
+      light_feature_norm = self.norm(self.light_feature['hidden_layer' +
+                                                        str(i)])
+      sim_middle_layer = tf.reduce_mean(
+          self.sim(booster_feature_no_gradient_norm, light_feature_norm))
+      return sim_middle_layer
+    else:
+      return tf.sqrt(
+          tf.reduce_sum(
+              tf.square(booster_feature_no_gradient -
+                        self.light_feature['hidden_layer' + str(i)])))
+
+  def build_predict_graph(self):
+    self.hidden_layer_feature_output = self._model_config.feature_based_distillation
+    if self._model_config.HasField('share_dnn'):
+      share_dnn_layer = dnn.DNN(self._model_config.share_dnn, self._l2_reg,
+                                'share_dnn', self._is_training)
+      share_feature = share_dnn_layer(self._features)
+    booster_dnn_layer = dnn.DNN(self._model_config.booster_dnn, self._l2_reg,
+                                'booster_dnn', self._is_training)
+    light_dnn_layer = dnn.DNN(self._model_config.light_dnn, self._l2_reg,
+                              'light_dnn', self._is_training)
+    if self._model_config.HasField('share_dnn'):
+      self.booster_feature = booster_dnn_layer(share_feature,
+                                               self.hidden_layer_feature_output)
+      input_embedding_stop_gradient = tf.stop_gradient(share_feature)
+      self.light_feature = light_dnn_layer(input_embedding_stop_gradient,
+                                           self.hidden_layer_feature_output)
+    else:
+      self.booster_feature = booster_dnn_layer(self._features,
+                                               self.hidden_layer_feature_output)
+      input_embedding_stop_gradient = tf.stop_gradient(self._features)
+      self.light_feature = light_dnn_layer(input_embedding_stop_gradient,
+                                           self.hidden_layer_feature_output)
+
+    if self._model_config.feature_based_distillation:
+      booster_out = tf.layers.dense(
+          self.booster_feature['hidden_layer_end'],
+          self._num_class,
+          kernel_regularizer=self._l2_reg,
+          name='booster_output')
+
+      light_out = tf.layers.dense(
+          self.light_feature['hidden_layer_end'],
+          self._num_class,
+          kernel_regularizer=self._l2_reg,
+          name='light_output')
+    else:
+      booster_out = tf.layers.dense(
+          self.booster_feature,
+          self._num_class,
+          kernel_regularizer=self._l2_reg,
+          name='booster_output')
+
+      light_out = tf.layers.dense(
+          self.light_feature,
+          self._num_class,
+          kernel_regularizer=self._l2_reg,
+          name='light_output')
+
+    self._prediction_dict.update(
+        self._output_to_prediction_impl(
+            booster_out,
+            self._loss_type,
+            num_class=self._num_class,
+            suffix='_booster'))
+    self._prediction_dict.update(
+        self._output_to_prediction_impl(
+            light_out,
+            self._loss_type,
+            num_class=self._num_class,
+            suffix='_light'))
+
+    return self._prediction_dict
+
+  def build_loss_graph(self):
+    logits_booster = self._prediction_dict['logits_booster']
+    logits_light = self._prediction_dict['logits_light']
+    self.feature_distillation_function = self._model_config.feature_distillation_function
+
+    # feature_based_distillation loss
+    if self._model_config.feature_based_distillation:
+      booster_hidden_units = self._model_config.booster_dnn.hidden_units
+      light_hidden_units = self._model_config.light_dnn.hidden_units
+      count = 0
+
+      for i, unit_i in enumerate(light_hidden_units):
+        for j, unit_j in enumerate(booster_hidden_units):
+          if light_hidden_units[i] == booster_hidden_units[j]:
+            self._prediction_dict[
+                'similarity_' + str(count)] = self.feature_based_sim(
+                    self._model_config.feature_based_distillation, i, j)
+            count += 1
+            break
+
+    self._loss_dict.update(
+        self._build_loss_impl(
+            LossType.CLASSIFICATION,
+            label_name=self._label_name,
+            loss_weight=self._sample_weight,
+            num_class=self._num_class,
+            suffix='_booster'))
+
+    self._loss_dict.update(
+        self._build_loss_impl(
+            LossType.CLASSIFICATION,
+            label_name=self._label_name,
+            loss_weight=self._sample_weight,
+            num_class=self._num_class,
+            suffix='_light'))
+
+    booster_logits_no_grad = tf.stop_gradient(logits_booster)
+
+    self._loss_dict['hint_loss'] = loss_builder.build(
+        LossType.L2_LOSS,
+        label=booster_logits_no_grad,
+        pred=logits_light,
+        loss_weight=self._sample_weight)
+
+    if self._model_config.feature_based_distillation:
+      for key, value in self._prediction_dict.items():
+        if key.startswith('similarity_'):
+          self._loss_dict[key] = -0.1 * value
+      return self._loss_dict
+    else:
+      return self._loss_dict
+
+  def build_metric_graph(self, eval_config):
+    metric_dict = {}
+    for metric in eval_config.metrics_set:
+      metric_dict.update(
+          self._build_metric_impl(
+              metric,
+              loss_type=LossType.CLASSIFICATION,
+              label_name=self._label_name,
+              num_class=self._num_class,
+              suffix='_light'))
+      metric_dict.update(
+          self._build_metric_impl(
+              metric,
+              loss_type=LossType.CLASSIFICATION,
+              label_name=self._label_name,
+              num_class=self._num_class,
+              suffix='_booster'))
+    return metric_dict
+
+  def get_outputs(self):
+    outputs = []
+    outputs.extend(
+        self._get_outputs_impl(
+            self._loss_type, self._num_class, suffix='_light'))
+    outputs.extend(
+        self._get_outputs_impl(
+            self._loss_type, self._num_class, suffix='_booster'))
+    return outputs