# -*- encoding:utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import logging import os import tensorflow as tf 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.protos.simi_pb2 import Similarity if tf.__version__ >= '2.0': tf = tf.compat.v1 losses = tf.losses class MatchModel(EasyRecModel): def __init__(self, model_config, feature_configs, features, labels=None, is_training=False): super(MatchModel, 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 if self._loss_type == LossType.CLASSIFICATION: assert self._num_class == 1 if self._loss_type in [LossType.CLASSIFICATION, LossType.L2_LOSS]: self._is_point_wise = True logging.info('Use point wise dssm.') else: self._is_point_wise = False logging.info('Use list wise dssm.') cls_mem = self._model_config.WhichOneof('model') sub_model_config = getattr(self._model_config, cls_mem) self._item_ids = None assert sub_model_config is not None, 'sub_model_config undefined: model_cls = %s' % cls_mem if getattr(sub_model_config, 'item_id', '') != '': logging.info('item_id feature is: %s' % sub_model_config.item_id) self._item_ids = features[sub_model_config.item_id] def _mask_in_batch(self, logits): batch_size = tf.shape(logits)[0] if getattr(self._model_config, 'ignore_in_batch_neg_sam', False): in_batch = logits[:, :batch_size] - ( 1 - tf.diag(tf.ones([batch_size], dtype=tf.float32))) * 1e32 return tf.concat([in_batch, logits[:, batch_size:]], axis=1) else: if self._item_ids is not None: mask_in_batch_neg = tf.to_float( tf.equal(self._item_ids[None, :batch_size], self._item_ids[:batch_size, None])) - tf.diag( tf.ones([batch_size], dtype=tf.float32)) tf.summary.scalar('in_batch_neg_conflict', tf.reduce_sum(mask_in_batch_neg)) return tf.concat([ logits[:, :batch_size] - mask_in_batch_neg * 1e32, logits[:, batch_size:]], axis=1) # yapf: disable else: return logits def _list_wise_sim(self, user_emb, item_emb): batch_size = tf.shape(user_emb)[0] hard_neg_indices = self._feature_dict.get('hard_neg_indices', None) if hard_neg_indices is not None: logging.info('With hard negative examples') noclk_size = tf.shape(hard_neg_indices)[0] # pos_item_emb, neg_item_emb, hard_neg_item_emb = tf.split( # item_emb, [batch_size, -1, noclk_size], axis=0) simple_item_emb, hard_neg_item_emb = tf.split( item_emb, [-1, noclk_size], axis=0) else: # pos_item_emb = item_emb[:batch_size] # neg_item_emb = item_emb[batch_size:] simple_item_emb = item_emb # pos_user_item_sim = tf.reduce_sum( # tf.multiply(user_emb, pos_item_emb), axis=1, keep_dims=True) # neg_user_item_sim = tf.matmul(user_emb, tf.transpose(neg_item_emb)) # simple_user_item_sim = tf.matmul(user_emb, tf.transpose(simple_item_emb)) _mode = os.environ['tf.estimator.mode'] if _mode == tf.estimator.ModeKeys.PREDICT: simple_user_item_sim = tf.reduce_sum( tf.multiply(user_emb, simple_item_emb), axis=1, keep_dims=True) else: simple_user_item_sim = tf.matmul(user_emb, tf.transpose(simple_item_emb)) if hard_neg_indices is None: return simple_user_item_sim else: user_emb_expand = tf.gather(user_emb, hard_neg_indices[:, 0]) hard_neg_user_item_sim = tf.reduce_sum( tf.multiply(user_emb_expand, hard_neg_item_emb), axis=1) max_num_neg = tf.reduce_max(hard_neg_indices[:, 1]) + 1 hard_neg_shape = tf.stack([tf.to_int64(batch_size), max_num_neg]) hard_neg_sim = tf.scatter_nd(hard_neg_indices, hard_neg_user_item_sim, hard_neg_shape) hard_neg_mask = tf.scatter_nd( hard_neg_indices, tf.ones_like(hard_neg_user_item_sim, dtype=tf.float32), shape=hard_neg_shape) # set tail positions to -1e32, so that after exp(x), will be zero hard_neg_user_item_sim = hard_neg_sim - (1 - hard_neg_mask) * 1e32 # user_item_sim = [pos_user_item_sim, neg_user_item_sim] # if hard_neg_indices is not None: # user_item_sim.append(hard_neg_user_item_sim) # return tf.concat(user_item_sim, axis=1) return tf.concat([simple_user_item_sim, hard_neg_user_item_sim], axis=1) def _point_wise_sim(self, user_emb, item_emb): user_item_sim = tf.reduce_sum( tf.multiply(user_emb, item_emb), axis=1, keep_dims=True) return user_item_sim def sim(self, user_emb, item_emb): # Name the outputs of the user tower and the item tower, i.e. the inputs of the # simularity operation. # Explicit names of these nodes are necessary for some online recall systems like # BasicEngine to split up the predicting graph into different clusters. user_emb = tf.identity(user_emb, 'user_tower_emb') item_emb = tf.identity(item_emb, 'item_tower_emb') if self._is_point_wise: return self._point_wise_sim(user_emb, item_emb) else: return self._list_wise_sim(user_emb, item_emb) def norm(self, fea): fea_norm = tf.nn.l2_normalize(fea, axis=-1) return fea_norm def build_predict_graph(self): if not self.has_backbone: raise NotImplementedError( 'method `build_predict_graph` must be implemented when you donot use backbone network' ) model = self._model_config.WhichOneof('model') assert model == 'model_params', '`model_params` must be configured' model_params = self._model_config.model_params for out in model_params.outputs: self._outputs.append(out) output = self.backbone user_tower_emb = output[model_params.user_tower_idx_in_output] item_tower_emb = output[model_params.item_tower_idx_in_output] if model_params.simi_func == Similarity.COSINE: user_tower_emb = self.norm(user_tower_emb) item_tower_emb = self.norm(item_tower_emb) temperature = model_params.temperature else: temperature = 1.0 user_item_sim = self.sim(user_tower_emb, item_tower_emb) / temperature if model_params.scale_simi: sim_w = tf.get_variable( 'sim_w', dtype=tf.float32, shape=(1), initializer=tf.ones_initializer()) sim_b = tf.get_variable( 'sim_b', dtype=tf.float32, shape=(1), initializer=tf.zeros_initializer()) y_pred = user_item_sim * tf.abs(sim_w) + sim_b else: y_pred = user_item_sim if self._is_point_wise: y_pred = tf.reshape(y_pred, [-1]) if self._loss_type == LossType.CLASSIFICATION: self._prediction_dict['logits'] = y_pred self._prediction_dict['probs'] = tf.nn.sigmoid(y_pred) elif self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY: y_pred = self._mask_in_batch(y_pred) self._prediction_dict['logits'] = y_pred self._prediction_dict['probs'] = tf.nn.softmax(y_pred) else: self._prediction_dict['y'] = y_pred self._prediction_dict['user_tower_emb'] = user_tower_emb self._prediction_dict['item_tower_emb'] = item_tower_emb self._prediction_dict['user_emb'] = tf.reduce_join( tf.as_string(user_tower_emb), axis=-1, separator=',') self._prediction_dict['item_emb'] = tf.reduce_join( tf.as_string(item_tower_emb), axis=-1, separator=',') return self._prediction_dict def build_loss_graph(self): if self._is_point_wise: return self._build_point_wise_loss_graph() else: return self._build_list_wise_loss_graph() def _build_list_wise_loss_graph(self): if self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY: batch_size = tf.shape(self._prediction_dict['probs'])[0] indices = tf.range(batch_size) indices = tf.concat([indices[:, None], indices[:, None]], axis=1) hit_prob = tf.gather_nd( self._prediction_dict['probs'][:batch_size, :batch_size], indices) sample_weights = tf.cast(tf.squeeze(self._sample_weight), tf.float32) self._loss_dict['cross_entropy_loss'] = -tf.reduce_mean( tf.log(hit_prob + 1e-12) * sample_weights) / tf.reduce_mean(sample_weights) logging.info('softmax cross entropy loss is used') user_features = self._prediction_dict['user_tower_emb'] pos_item_emb = self._prediction_dict['item_tower_emb'][:batch_size] pos_simi = tf.reduce_sum(user_features * pos_item_emb, axis=1) # if pos_simi < 0, produce loss reg_pos_loss = tf.nn.relu(-pos_simi) self._loss_dict['reg_pos_loss'] = tf.reduce_mean( reg_pos_loss * sample_weights) / tf.reduce_mean(sample_weights) # the AMM loss for DAT model if all([ k in self._prediction_dict.keys() for k in ['augmented_p_u', 'augmented_p_i', 'augmented_a_u', 'augmented_a_i'] ]): self._loss_dict[ 'amm_loss_u'] = self._model_config.amm_u_weight * tf.reduce_mean( tf.square(self._prediction_dict['augmented_a_u'] - self._prediction_dict['augmented_p_i'][:batch_size]) * sample_weights) / tf.reduce_mean(sample_weights) self._loss_dict[ 'amm_loss_i'] = self._model_config.amm_i_weight * tf.reduce_mean( tf.square(self._prediction_dict['augmented_a_i'][:batch_size] - self._prediction_dict['augmented_p_u']) * sample_weights) / tf.reduce_mean(sample_weights) else: raise ValueError('invalid loss type: %s' % str(self._loss_type)) return self._loss_dict def _build_point_wise_loss_graph(self): label = list(self._labels.values())[0] if self._loss_type == LossType.CLASSIFICATION: pred = self._prediction_dict['logits'] loss_name = 'cross_entropy_loss' elif self._loss_type == LossType.L2_LOSS: pred = self._prediction_dict['y'] loss_name = 'l2_loss' else: raise ValueError('invalid loss type: %s' % str(self._loss_type)) kwargs = {'loss_name': loss_name} self._loss_dict[loss_name] = loss_builder.build( self._loss_type, label=label, pred=pred, loss_weight=self._sample_weight, **kwargs) # 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_graph(self, eval_config): if self._is_point_wise: return self._build_point_wise_metric_graph(eval_config) else: return self._build_list_wise_metric_graph(eval_config) def _build_list_wise_metric_graph(self, eval_config): from easy_rec.python.core.easyrec_metrics import metrics_tf logits = self._prediction_dict['logits'] # label = tf.zeros_like(logits[:, :1], dtype=tf.int64) batch_size = tf.shape(logits)[0] label = tf.cast(tf.range(batch_size), tf.int64) indices = tf.range(batch_size) indices = tf.concat([indices[:, None], indices[:, None]], axis=1) pos_item_sim = tf.gather_nd(logits[:batch_size, :batch_size], indices) metric_dict = {} for metric in eval_config.metrics_set: if metric.WhichOneof('metric') == 'recall_at_topk': metric_dict['recall@%d' % metric.recall_at_topk.topk] = metrics_tf.recall_at_k( label, logits, metric.recall_at_topk.topk) logits_v2 = tf.concat([pos_item_sim[:, None], logits[:, batch_size:]], axis=1) labels_v2 = tf.zeros_like(logits_v2[:, :1], dtype=tf.int64) metric_dict['recall_neg_sam@%d' % metric.recall_at_topk.topk] = metrics_tf.recall_at_k( labels_v2, logits_v2, metric.recall_at_topk.topk) metric_dict['recall_in_batch@%d' % metric.recall_at_topk.topk] = metrics_tf.recall_at_k( label, logits[:, :batch_size], metric.recall_at_topk.topk) else: raise ValueError('invalid metric type: %s' % str(metric)) return metric_dict def _build_point_wise_metric_graph(self, eval_config): from easy_rec.python.core.easyrec_metrics import metrics_tf metric_dict = {} label = list(self._labels.values())[0] for metric in eval_config.metrics_set: if metric.WhichOneof('metric') == 'auc': assert self._loss_type == LossType.CLASSIFICATION metric_dict['auc'] = metrics_tf.auc(label, self._prediction_dict['probs']) elif metric.WhichOneof('metric') == 'mean_absolute_error': assert self._loss_type == LossType.L2_LOSS metric_dict['mean_absolute_error'] = metrics_tf.mean_absolute_error( tf.to_float(label), self._prediction_dict['y']) else: raise ValueError('invalid metric type: %s' % str(metric)) return metric_dict def get_outputs(self): if not self.has_backbone: raise NotImplementedError( 'could not call get_outputs on abstract class MatchModel') if self._loss_type == LossType.CLASSIFICATION: return [ 'logits', 'probs', 'user_emb', 'item_emb', 'user_tower_emb', 'item_tower_emb' ] elif self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY: self._prediction_dict['logits'] = tf.squeeze( self._prediction_dict['logits'], axis=-1) self._prediction_dict['probs'] = tf.nn.sigmoid( self._prediction_dict['logits']) return [ 'logits', 'probs', 'user_emb', 'item_emb', 'user_tower_emb', 'item_tower_emb' ] elif self._loss_type == LossType.L2_LOSS: return ['y', 'user_emb', 'item_emb', 'user_tower_emb', 'item_tower_emb'] else: raise ValueError('invalid loss type: %s' % str(self._loss_type))