def main(_): """ 数据输入部分 """ # 数据集路径 path = "data/criteo_sample.txt" # 加载数据集 sparse_features, dense_features, sparse_feature_columns, dense_feature_columns, label = input_util.load_criteo_data( path, sep='\t') # 划分训练集 & 验证集 train_size = int(len(label) * 0.85) train_sparse_features = sparse_features[:train_size] train_dense_features = dense_features[:train_size] train_label = label[:train_size] eval_sparse_features = sparse_features[train_size:] eval_dense_features = dense_features[train_size:] eval_label = label[train_size:] # 指定 wide & deep 模型的 wide部分和deep部分所需要的feature linear_feature_names = [each.name for each in sparse_feature_columns ] + [each.name for each in dense_feature_columns] dnn_feature_names = [each.name for each in sparse_feature_columns] # 处理model params model_params = { "sparse_feature_columns": sparse_feature_columns, "batch_size": FLAGS.batch_size, "embedding_size": 8, "attention_size": 4, "lr": 0.01 } # 清空model dir model_util.clean_model(FLAGS.model_dir) config = tf.estimator.RunConfig() estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.model_dir, params=model_params, config=config) if FLAGS.train == 1: estimator.train(input_fn=lambda: input_fn(train_sparse_features, train_dense_features, train_label, epoch=FLAGS.epoch, batch_size=FLAGS.batch_size)) if FLAGS.eval == 1: estimator.evaluate( input_fn=lambda: input_fn(eval_sparse_features, eval_dense_features, eval_label, epoch=FLAGS.epoch, batch_size=FLAGS.batch_size))
def main(_): """ 数据输入部分 """ # 数据集路径 path = "data/census-income.data.gz" # 加载数据集 sparse_features, dense_features, sparse_feature_columns, dense_feature_columns, label = input_util.load_census_income_data( path, sep=',') # 划分训练集 & 验证集 train_size = int(len(label) * 0.85) train_sparse_features = sparse_features[:train_size] train_dense_features = dense_features[:train_size] train_label = label[:train_size] eval_sparse_features = sparse_features[train_size:] eval_dense_features = dense_features[train_size:] eval_label = label[train_size:] # 处理model params model_params = { "sparse_feature_columns": sparse_feature_columns, "dense_feature_columns": dense_feature_columns, "batch_size": FLAGS.batch_size, "embedding_size": 8, "task_nums": 2, "lr": 0.01 } # 清空model dir model_util.clean_model(FLAGS.model_dir) config = tf.estimator.RunConfig() estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.model_dir, params=model_params, config=config) if FLAGS.train == 1: estimator.train(input_fn=lambda: input_fn(train_sparse_features, train_dense_features, train_label, epoch=FLAGS.epoch, batch_size=FLAGS.batch_size, if_multi_label=True)) if FLAGS.eval == 1: estimator.evaluate( input_fn=lambda: input_fn(eval_sparse_features, eval_dense_features, eval_label, epoch=FLAGS.epoch, batch_size=FLAGS.batch_size, if_multi_label=True))