devbox
Cpp五条
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

利用tensorflow estimator API实现双塔推荐算法

作者：Cpp五条 | 2024-02-21 12:08:45

赞

踩

利用tensorflow estimator API实现双塔推荐算法

本文完整代码见： https://github.com/cdj0311/two_tower_recommendation_system

Tensorflow estimator实现分布式训练很简单，只需要将数据进行相应的切分丢给模型就可以很方便的完成分布式训练了。以下代码是一个完整的推荐算法模板，可根据自己的需要修改数据读取和模型结构部分，tensorflow==1.13.1。

1. 特征处理部分，feature_processing.py

  1. #coding:utf-8
  2. import tensorflow as tf
  3. from tensorflow import feature_column as fc
  4. import config
  5.  
  6. FLAGS = config.FLAGS
  7.  
  8. class FeatureConfig(object):
  9.     def __init__(self):
  10.         self.user_columns = dict()
  11.         self.item_columns = dict()
  12.         self.feature_spec = dict()
  13.  
  14.     def create_features_columns(self):
  15.         # 向量类特征
  16.         user_vector = fc.numeric_column(key="user_vector", shape=(128,), default_value=[0.0] * 128, dtype=tf.float32)
  17.         item_vector = fc.numeric_column(key="item_vector", shape=(128,), default_value=[0.0] * 128, dtype=tf.float32)
  18.         
  19.         # 分桶类特征
  20.         age = fc.numeric_column(key="age", shape=(1,), default_value=[0], dtype=tf.int64)
  21.         age = fc.bucketized_column(input_fc, boundaries=[0,10,20,30,40,50,60,70,80])
  22.         age = fc.embedding_column(age, dimension=32, combiner='mean')
  23.         
  24.         # 分类特征
  25.         city = fc.categorical_column_with_identity(key="city", num_buckets=1000, default_value=0)
  26.         city = fc.embedding_column(city, dimension=32, combiner='mean')
  27.         
  28.         # hash特征
  29.         device_id = fc.categorical_column_with_hash_bucket(key="device_id", 
  30.                     hash_bucket_size=1000000, dtype=tf.int64)
  31.         device_id = fc.embedding_column(device_id, dimension=32, combiner='mean')
  32.  
  33.         item_id = fc.categorical_column_with_hash_bucket(key="item_id", 
  34.                     hash_bucket_size=1000000, dtype=tf.int64)
  35.         item_id = fc.embedding_column(item_id, dimension=32, combiner='mean')
  36.         
  37.         self.user_columns["user_vector"] = user_vector
  38.         self.user_columns["age"] = age
  39.         self.user_columns["city"] = city
  40.         self.user_columns["device_id"] = device_id
  41.         self.item_columns["item_vector"] = item_vector
  42.         self.item_columns["item_id"] = item_id
  43.  
  44.         self.feature_spec = tf.feature_column.make_parse_example_spec(self.user_columns.values()+self.item_columns.values())
  45.  
  46.         return self

2. 模型部分，base_model.py

  1. # coding:utf-8
  2. import tensorflow as tf
  3. from tensorflow import feature_column as fc
  4. import config
  5. import feature_processing as fe
  6.  
  7. FLAGS = config.FLAGS
  8.  
  9.  
  10. def build_user_model(features, mode, params):
  11.     # 特征输入
  12.     feature_inputs = []
  13.     for key, value in params["feature_configs"].user_columns.items():
  14.         feature_inputs.append(tf.input_layer(features, value))
  15.     # 特征拼接
  16.     net = tf.concat(feature_inputs, axis=1)
  17.     # 全连接
  18.     for idx, units in enumerate(params["hidden_units"]):
  19.         net = tf.layers.dense(net, units=units, activation=tf.nn.leaky_relu, name="fc_layer_%s"%idx)
  20.         net = tf.layers.dropout(net, 0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
  21.     # 最后输出
  22.     net = tf.layers.dense(net, units=128, name="user_output_layer")
  23.     return net
  24.  
  25. def build_item_model(features, mode, params):
  26.     # 特征输入
  27.     feature_inputs = []
  28.     for key, value in params["feature_configs"].item_columns.items():
  29.         feature_inputs.append(tf.input_layer(features, value))
  30.     # 特征拼接
  31.     net = tf.concat(feature_inputs, axis=1)
  32.     # 全连接
  33.     for idx, units in enumerate(params["hidden_units"]):
  34.         net = tf.layers.dense(net, units=units, activation=tf.nn.leaky_relu, name="fc_layer_%s"%idx)
  35.         net = tf.layers.dropout(net, 0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
  36.     # 最后输出
  37.     net = tf.layers.dense(net, units=128, name="item_output_layer")
  38.     return net
  39.  
  40. def model_fn(features, labels, mode, params):
  41.     user_net = build_user_model(features, mode, params)
  42.     item_net = build_item_model(features, mode, params)
  43.     dot = tf.reduce_sum(tf.multiply(user_net, item_net), axis=1, keepdims=True)
  44.     pred = tf.sigmoid(dot)
  45.     if mode == tf.estimator.ModeKeys.PREDICT:
  46.         return tf.estimator.EstimatorSpec(mode, predictions={"output": pred})
  47.     if mode == tf.estimator.ModeKeys.EVAL:
  48.         loss = tf.losses.log_loss(labels, pred)
  49.         metrics = {"auc": tf.metrics.auc(labels=labels, predictions=pred)}
  50.         return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
  51.     if mode == tf.estimator.ModeKeys.TRAIN:
  52.         loss = tf.losses.log_loss(labels, pred)
  53.         global_step = tf.train.get_global_step()
  54.         learning_rate = tf.train.exponential_decay(params["learning_rate"], global_step, 100000, 0.9, staircase=True)
  55.         train_op = (tf.train.AdagradOptimizer(learning_rate).minimize(loss, global_step=global_step))
  56.         return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

3. 数据输入部分，data_input.py

  1. # coding:utf-8
  2. import tensorflow as tf
  3. import config
  4.  
  5. FLAGS = config.FLAGS
  6.  
  7. def parse_exp(example):
  8.     features_def = {}
  9.     features_def["label"] = tf.io.FixedLenFeature([1], tf.int64)
  10.     """
  11.     数据解析逻辑
  12.     """
  13.     features = tf.io.parse_single_example(example, features_def)
  14.     label = features["label"]
  15.     return features, label
  16.  
  17.  
  18. def train_input_fn(filenames=None, batch_size=128):
  19.     with tf.gfile.Open(filenames) as f:
  20.         filenames = f.read().split()
  21.     files_all = []
  22.     for f in filenames:
  23.         files_all += tf.gfile.Glob(f)
  24.     train_worker_num = len(FLAGS.worker_hosts.split(","))
  25.     hash_id = FLAGS.task_index if FLAGS.job_name == "worker" else train_worker_num - 1
  26.     files_shard = [files for i, files in enumerate(files_all) if i % train_worker_num == hash_id]
  27.     dataset = tf.data.TFRecordDataset(files_shard)
  28.     dataset = dataset.shuffle(batch_size*100)
  29.     dataset = dataset.map(parse_exp, num_parallel_calls=8)
  30.     dataset = dataset.repeat().batch(batch_size).prefetch(1)
  31.     return dataset
  32.  
  33. def eval_input_fn(filenames=None, batch_size=128):
  34.     with tf.gfile.Open(filenames) as f:
  35.         filenames = f.read().split()
  36.     files = tf.data.Dataset.list_files(filenames)
  37.     dataset = files.apply(tf.contrib.data.parallel_interleave(lambda filename: tf.data.TFRecordDataset(files), buffer_output_elements=batch_size*20, cycle_length=10))
  38.     dataset = dataset.map(parse_exp, num_parallel_calls=4)
  39.     dataset = dataset.batch(batch_size)
  40.     return dataset

4. 主函数部分，main.py

  1. # encoding:utf-8
  2. import os
  3. import tensorflow as tf
  4. from tensorflow import feature_column as fc
  5. import feature_engineering as fe
  6. from feature_processing import FeatureConfig
  7. import base_model
  8. import data_input
  9. import config
  10.  
  11. FLAGS = config.FLAGS
  12.  
  13. if FLAGS.run_on_cluster:
  14.     cluster = json.loads(os.environ["TF_CONFIG"])
  15.     task_index = int(os.environ["TF_INDEX"])
  16.     task_type = os.environ["TF_ROLE"]
  17.  
  18.  
  19. def main(unused_argv):
  20.     feature_configs = FeatureConfig().create_features_columns()
  21.     classifier = tf.estimator.Estimator(model_fn=model.model_fn,
  22.                                         config=tf.estimator.RunConfig(model_dir=FLAGS.model_dir,
  23.                                                                       save_checkpoints_steps=FLAGS.save_checkpoints_steps,
  24.                                                                       keep_checkpoint_max=3),
  25.                                         params={"feature_configs": feature_configs,
  26.                                                 "hidden_units": list(map(int, FLAGS.hidden_units.split(","))),
  27.                                                 "learning_rate": FLAGS.learning_rate}
  28.                                         )
  29.     def train_eval_model():
  30.         train_spec = tf.estimator.TrainSpec(input_fn=lambda: data_input.train_input_fn(FLAGS.train_data, FLAGS.batch_size),
  31.                                             max_steps=FLAGS.train_steps)
  32.         eval_spec = tf.estimator.EvalSpec(input_fn=lambda: data_input.eval_input_fn(FLAGS.eval_data, FLAGS.batch_size),
  33.                                           start_delay_secs=60,
  34.                                           throttle_secs = 30,
  35.                                           steps=1000)
  36.         tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)
  37.  
  38.     def export_model():
  39.         feature_spec = feature_configs.feature_spec
  40.         feature_map = {}
  41.         for key, feature in feature_spec.items():
  42.             if key not in fe.feature_configs:
  43.                 continue
  44.             if isinstance(feature, tf.io.VarLenFeature):  # 可变长度
  45.                 feature_map[key] = tf.placeholder(dtype=feature.dtype, shape=[1], name=key)
  46.             elif isinstance(feature, tf.io.FixedLenFeature):  # 固定长度
  47.                 feature_map[key] = tf.placeholder(dtype=feature.dtype, shape=[None, feature.shape[0]], name=key)
  48.         serving_input_recevier_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(feature_map)
  49.         export_dir = classifier.export_saved_model(FLAGS.output_model, serving_input_recevier_fn)
  50.  
  51.     # 模型训练
  52.     train_eval_model()
  53.     
  54.     # 导出模型，只在chief中导出一次即可
  55.     if FLAGS.run_on_cluster: 
  56.         if task_type == "chief":
  57.             export_model()
  58.     else:
  59.         export_model()
  60.  
  61.  
  62. if __name__ == "__main__":
  63.     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
  64.     tf.logging.set_verbosity(tf.logging.INFO)
  65.     tf.app.run(main=main)

5. config.py

  1. import json, os, re, codecs
  2. import tensorflow as tf
  3.  
  4. flags = tf.app.flags
  5.  
  6. flags.DEFINE_boolean("run_on_cluster", True, "Whether the cluster info need to be passed in as input")
  7.  
  8. flags.DEFINE_string("train_dir", "", "")
  9. flags.DEFINE_string("data_dir", "", "")
  10. flags.DEFINE_string("log_dir", "", "")
  11. flags.DEFINE_string("ps_hosts", "","")
  12. flags.DEFINE_string("worker_hosts", "","")
  13. flags.DEFINE_string("job_name", "", "One of 'ps', 'worker'")
  14. flags.DEFINE_integer("task_index", 0, "Index of task within the job")
  15. flags.DEFINE_string("model_dir", "hdfs:/user/ranker/ckpt/", "Base directory for the model.")
  16. flags.DEFINE_string("output_model", "hdfs:/user/ranker/model/", "Saved model.")
  17. flags.DEFINE_string("train_data", "./train_data.txt", "Directory for storing mnist data")
  18. flags.DEFINE_string("eval_data", "./eval_data.txt", "Path to the evaluation data.")
  19. flags.DEFINE_string("hidden_units", "512,256,128", "hidden units.")
  20. flags.DEFINE_integer("train_steps",1000000, "Number of (global) training steps to perform")
  21. flags.DEFINE_integer("batch_size", 256, "Training batch size")
  22. flags.DEFINE_float("learning_rate", 0.0001, "learning rate")
  23. flags.DEFINE_integer("save_checkpoints_steps", 10000, "Save checkpoints every this many steps")
  24.  
  25. FLAGS = flags.FLAGS

6. 最后是任务提交命令，各公司应该都有自己的提交平台，这里就不贴出来了。

当然该代码直接单机运行也是没有任何问题的，只需要将run_on_cluster设置为False即可。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Cpp五条/article/detail/123638
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号