谷歌双塔模型原理及Estimator实现

背景介绍

• 文章核心?
  目前在大数据推荐系统环境下, 基于模型的建模方式有很多,利用双塔模型对User和Item建模是最经典的模型之一.双塔模型通过独立训练用户特征和文章特征,通过Embedding的方式去表达用户和物品的属性,在线上通过向量计算得到用户对物品的Label信息.

• 应用场景
  推荐系统的两个场景包括召回和推荐.
  召回
  召回的作用是从千万级别的海量数据中找出千百级别左右的用户可能喜欢的物品,主要目的是解决海量数据处理;
  对于召回来说: 双塔可以根据用户和物品的Embedding表示,去做向量计算,这样有利于海量数据的计算.
  排序
  排序的作用是从召回的千百级别的物品对用户做个性化精准排序,主要目的是考虑更多用户特征,比如点击序列特征,通过更复杂的模型学习用户属性.
  对于排序来说: 双塔可以考虑用户和物品更多的特征,以端到端的方式计算用户对物品的Label结果

模型结构

这里简单说几点:

• 用户和物品的网络可以根据自己需求来改相应的结构;
• DSSM模型结构跟双塔类似,塔的数量可能不一样;
• 线上可以用faiss向量索引工具计算向量之间的相似性,提升计算效率

Estimator实现

这块是通过tensorflow的高阶api去构建用户侧和物品侧的网络,一个简单的DNN网络demo,大家可以根据需求更改自己的网络结构.

#!/usr/bin/env python3

#############################################
# FileName: model.py
# Descreption: 双塔召回模型
#############################################

import tensorflow as tf
import config
# 配置参数文件
FLAGS = config.FLAGS

def build_user_model(features, mode, params):
    """
    User Tower Feature Embedding
    :param features:
    :param mode:
    :param params:
    :return:
    """
    # 特征输入
    user_vector_net = []
    for key, value in params["feature_configs"].user_columns.items():
        user_vector_net.append(tf.feature_column.input_layer(features, value))
    # 特征拼接
    net = tf.concat(user_vector_net, axis=1)
    # 全连接
    for idx, units in enumerate(params["hidden_units"]):
        net = tf.layers.dense(net, units=units, activation=tf.nn.leaky_relu, name="user_fc_layer_%s"%idx)
        net = tf.layers.dropout(net, 0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
    # 最后输出
    net = tf.layers.dense(net, units=64, name="user_output_layer")
    return net


def build_item_model(features, mode, params):
    """
    Item Tower Feature Embedding
    :param features:
    :param model:
    :param params:
    :return:
    """
    # 特征输入
    item_vector_net = []
    for key, value in params["feature_configs"].items_columns.items():
        item_vector_net.append(tf.feature_column.input_layer(features, value))
    # 特征拼接
    net = tf.concat(item_vector_net, axis=1)
    # 全连接
    for idx, units in enumerate(params["hidden_units"]):
        net = tf.layers.dense(net, units=units, activation=tf.nn.leaky_relu, name="user_fc_layer_%s"%idx)
        net = tf.layers.dropout(net, 0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
    # 最后输出
    net = tf.layers.dense(net, units=64, name="item_output_layer")
    return net


def model_fn(features, labels, mode, params):
    # predict
    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = dict()
        if FLAGS.export_user_model:
            user_net = build_user_model(features, mode, params)
            predictions = {"user_vector": user_net}
        elif FLAGS.export_item_model:
            item_net = build_item_model(features, mode, params)
            predictions = {"item_vector": item_net}
        export_outputs = {"prediction": tf.estimator.export.PredictOutput(outputs=predictions)}
        return tf.estimator.EstimatorSpec(mode, predictions=predictions, export_outputs=export_outputs)

    # 构建user侧模型
    user_net = build_user_model(features, mode, params)
    # 构建item侧模型
    item_net = build_item_model(features, mode, params)

    # 向量计算
    dot = tf.reduce_sum(tf.multiply(user_net, item_net), axis=1, keep_dims=True)
    pred = tf.sigmoid(dot)

    if mode == tf.estimator.ModeKeys.EVAL:
        labels = tf.cast(labels, tf.float32)
        loss = tf.losses.log_loss(labels, pred)
        metrics = {"auc": tf.metrics.auc(labels=labels, predictions=pred)}
        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)

    if mode == tf.estimator.ModeKeys.TRAIN:
        labels = tf.cast(labels, tf.float32)
        loss = tf.losses.log_loss(labels, pred)
        global_step = tf.train.get_global_step()
        train_op = tf.train.AdagradOptimizer(FLAGS.learning_rate).minimize(loss, global_step=global_step)
        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

其他的模块都是estimator的公共模块,根据自己的需求来改就可以了.
总的来说: 双塔模型算是深度学习中比较简单的,容易实现的模型了!