修改测试文件

tensnsorflow/train_multi.py

@@ -8,10 +8,7 @@
 import shutil
 import os
 import json
-import glob
 from datetime import date, timedelta
-
-import random
 import tensorflow as tf
 
 #################### CMD Arguments ####################
@@ -37,7 +34,8 @@ tf.app.flags.DEFINE_string("deep_layers", '256,128,64', "deep layers")
 tf.app.flags.DEFINE_string("dropout", '0.5,0.5,0.5', "dropout rate")
 tf.app.flags.DEFINE_boolean("batch_norm", False, "perform batch normaization (True or False)")
 tf.app.flags.DEFINE_float("batch_norm_decay", 0.9, "decay for the moving average(recommend trying decay=0.9)")
-tf.app.flags.DEFINE_string("data_dir", '', "data dir")
+tf.app.flags.DEFINE_string("hdfs_dir", '', "hdfs dir")
+tf.app.flags.DEFINE_string("local_dir", '', "local dir")
 tf.app.flags.DEFINE_string("dt_dir", '', "data dt partition")
 tf.app.flags.DEFINE_string("model_dir", '', "model check point dir")
 tf.app.flags.DEFINE_string("servable_model_dir", '', "export servable model for TensorFlow Serving")
@@ -301,7 +299,8 @@ def main(_):
 
     print('task_type ', FLAGS.task_type)
     print('model_dir ', FLAGS.model_dir)
-    print('data_dir ', FLAGS.data_dir)
+    print('hdfs_dir ', FLAGS.hdfs_dir)
+    print('local_dir ', FLAGS.local_dir)
     print('dt_dir ', FLAGS.dt_dir)
     print('num_epochs ', FLAGS.num_epochs)
     print('feature_size ', FLAGS.feature_size)
@@ -320,6 +319,7 @@ def main(_):
     path = "hdfs:///strategy/esmm/"
     tr_files = [path+"tr/part-r-00000"]
     va_files = [path+"va/part-r-00000"]
+    te_files = ["%s/part-r-00000" % FLAGS.hdfs_dir]
 
     # tr_files = glob.glob("%s/tr/*tfrecord" % FLAGS.data_dir)
     # random.shuffle(tr_files)
@@ -366,9 +366,9 @@ def main(_):
             print('%s: %s' % (key,value))
     elif FLAGS.task_type == 'infer':
         preds = Estimator.predict(input_fn=lambda: input_fn(te_files, num_epochs=1, batch_size=FLAGS.batch_size), predict_keys=["pctcvr","pctr","pcvr"])
-        with open(FLAGS.data_dir+"/pred.txt", "w") as fo:
+        with open(FLAGS.local_dir+"/pred.txt", "w") as fo:
             print("-"*100)
-            with open(FLAGS.data_dir + "/pred.txt", "w") as fo:
+            with open(FLAGS.local_dir + "/pred.txt", "w") as fo:
                 for prob in preds:
                     fo.write("%f\t%f\t%f\n" % (prob['pctr'], prob['pcvr'], prob['pctcvr']))
     elif FLAGS.task_type == 'export':