try predict

src/main.py

@@ -29,7 +29,7 @@ def main():
 
     all_features = build_features(df)
 
-    params = {"feature_columns": all_features, "hidden_units": [32], "learning_rate": 0.1}
+    params = {"feature_columns": all_features, "hidden_units": [64, 32], "learning_rate": 0.1}
     model_path = str(Path("~/data/model_tmp/").expanduser())
     if os.path.exists(model_path):
         shutil.rmtree(model_path)
@@ -48,10 +48,7 @@ def main():
 
     for i in range(10):
         test_300 = test_df.sample(300)
-        time_1 = timeit.default_timer()
         model_predict(test_300, predict_fn)
-        total_1 = (timeit.default_timer() - time_1)
-        print("prediction cost {:.5f} s".format(total_1))
 
     total_time = (time.time() - time_begin) / 60
     print("cost {:.2f} mins at {}".format(total_time, datetime.now()))

src/models/esmm/model.py

@@ -2,6 +2,7 @@ import tensorflow as tf
 from tensorflow import feature_column as fc
 from tensorflow.python.estimator.canned import head as head_lib
 from tensorflow.python.ops.losses import losses
+import timeit
 
 
 def build_deep_layer(net, params):
@@ -90,7 +91,6 @@ def _bytes_feature(value):
 
 
 def model_predict(inputs, predict_fn):
-
     int_columns = [
         "active_type", "active_days", "card_id", "is_pure_author", "is_have_reply", "is_have_pure_reply", "content_level",
         "topic_num", "favor_num", "vote_num"
@@ -110,6 +110,9 @@ def model_predict(inputs, predict_fn):
                 features[col] = _bytes_feature(str(value).encode(encoding="utf-8"))
         example = tf.train.Example(features=tf.train.Features(feature=features))
         examples.append(example.SerializeToString())
+    time_1 = timeit.default_timer()
     predictions = predict_fn({"examples": examples})
+    total_1 = (timeit.default_timer() - time_1)
+    print("prediction cost {:.5f} s".format(total_1))
     # print(predictions)
     return predictions