# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import tensorflow as tf
from easy_rec.python.layers import dnn
from easy_rec.python.model.match_model import MatchModel
from easy_rec.python.protos.dssm_pb2 import DSSM as DSSMConfig
from easy_rec.python.protos.loss_pb2 import LossType
from easy_rec.python.protos.simi_pb2 import Similarity
from easy_rec.python.utils.proto_util import copy_obj
if tf.__version__ >= '2.0':
tf = tf.compat.v1
losses = tf.losses
[docs]class DSSM(MatchModel):
[docs] def __init__(self,
model_config,
feature_configs,
features,
labels=None,
is_training=False):
super(DSSM, self).__init__(model_config, feature_configs, features, labels,
is_training)
assert self._model_config.WhichOneof('model') == 'dssm', \
'invalid model config: %s' % self._model_config.WhichOneof('model')
self._model_config = self._model_config.dssm
assert isinstance(self._model_config, DSSMConfig)
# copy_obj so that any modification will not affect original config
self.user_tower = copy_obj(self._model_config.user_tower)
self.user_tower_feature, _ = self._input_layer(self._feature_dict, 'user')
# copy_obj so that any modification will not affect original config
self.item_tower = copy_obj(self._model_config.item_tower)
self.item_tower_feature, _ = self._input_layer(self._feature_dict, 'item')
self._user_tower_emb = None
self._item_tower_emb = None
[docs] def build_predict_graph(self):
num_user_dnn_layer = len(self.user_tower.dnn.hidden_units)
last_user_hidden = self.user_tower.dnn.hidden_units.pop()
user_dnn = dnn.DNN(self.user_tower.dnn, self._l2_reg, 'user_dnn',
self._is_training)
user_tower_emb = user_dnn(self.user_tower_feature)
user_tower_emb = tf.layers.dense(
inputs=user_tower_emb,
units=last_user_hidden,
kernel_regularizer=self._l2_reg,
name='user_dnn/dnn_%d' % (num_user_dnn_layer - 1))
num_item_dnn_layer = len(self.item_tower.dnn.hidden_units)
last_item_hidden = self.item_tower.dnn.hidden_units.pop()
item_dnn = dnn.DNN(self.item_tower.dnn, self._l2_reg, 'item_dnn',
self._is_training)
item_tower_emb = item_dnn(self.item_tower_feature)
item_tower_emb = tf.layers.dense(
inputs=item_tower_emb,
units=last_item_hidden,
kernel_regularizer=self._l2_reg,
name='item_dnn/dnn_%d' % (num_item_dnn_layer - 1))
if self._loss_type == LossType.CLASSIFICATION:
if self._model_config.simi_func == Similarity.COSINE:
user_tower_emb = self.norm(user_tower_emb)
item_tower_emb = self.norm(item_tower_emb)
user_item_sim = self.sim(user_tower_emb, item_tower_emb)
if self._model_config.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
[docs] def get_outputs(self):
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))
[docs] def build_output_dict(self):
output_dict = super(DSSM, self).build_output_dict()
output_dict['user_tower_feature'] = tf.reduce_join(
tf.as_string(self.user_tower_feature), axis=-1, separator=',')
output_dict['item_tower_feature'] = tf.reduce_join(
tf.as_string(self.item_tower_feature), axis=-1, separator=',')
return output_dict
[docs] def build_rtp_output_dict(self):
output_dict = super(DSSM, self).build_rtp_output_dict()
if 'user_tower_emb' not in self._prediction_dict:
raise ValueError(
'User tower embedding does not exist. Please checking predict graph.')
output_dict['user_embedding_output'] = tf.identity(
self._prediction_dict['user_tower_emb'], name='user_embedding_output')
if 'item_tower_emb' not in self._prediction_dict:
raise ValueError(
'Item tower embedding does not exist. Please checking predict graph.')
output_dict['item_embedding_output'] = tf.identity(
self._prediction_dict['item_tower_emb'], name='item_embedding_output')
if self._loss_type == LossType.CLASSIFICATION:
if 'probs' not in self._prediction_dict:
raise ValueError(
'Probs output does not exist. Please checking predict graph.')
output_dict['rank_predict'] = tf.identity(
self._prediction_dict['probs'], name='rank_predict')
return output_dict