update steps

src/main.py

@@ -42,7 +42,7 @@ def main():
     #     shutil.rmtree(model_path)
 
     model = tf.estimator.Estimator(model_fn=esmm_model_fn, params=params, model_dir=model_path)
-    train_spec = tf.estimator.TrainSpec(input_fn=lambda: esmm_input_fn(train_df, shuffle=True), max_steps=None)
+    train_spec = tf.estimator.TrainSpec(input_fn=lambda: esmm_input_fn(train_df, shuffle=True), max_steps=10000)
     eval_spec = tf.estimator.EvalSpec(input_fn=lambda: esmm_input_fn(val_df, shuffle=False))
     tf.estimator.train_and_evaluate(model, train_spec, eval_spec)
 

src/models/esmm/model.py

@@ -58,9 +58,7 @@ def esmm_model_fn(features, labels, mode, params):
         ctcvr_loss = tf.reduce_sum(tf.compat.v1.losses.log_loss(labels=cvr_labels, predictions=ctcvr_preds))
         loss = ctr_loss + ctcvr_loss
 
-        if mode == tf.estimator.ModeKeys.EVAL:
-            ctr_accuracy = tf.compat.v1.metrics.accuracy(labels=ctr_labels,
-                                                         predictions=tf.to_float(tf.greater_equal(ctr_preds, 0.5)))
+        ctr_accuracy = tf.compat.v1.metrics.accuracy(labels=ctr_labels, predictions=tf.to_float(tf.greater_equal(ctr_preds, 0.5)))
         ctcvr_accuracy = tf.compat.v1.metrics.accuracy(labels=cvr_labels,
                                                        predictions=tf.to_float(tf.greater_equal(ctcvr_preds, 0.5)))
         ctr_auc = tf.compat.v1.metrics.auc(labels=ctr_labels, predictions=ctr_preds)
@@ -70,9 +68,10 @@ def esmm_model_fn(features, labels, mode, params):
         tf.compat.v1.summary.scalar("ctcvr_accuracy", ctcvr_accuracy[1])
         tf.compat.v1.summary.scalar("ctr_auc", ctr_auc[1])
         tf.compat.v1.summary.scalar("ctcvr_auc", ctcvr_auc[1])
+        if mode == tf.estimator.ModeKeys.EVAL:
             return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
         train_op = optimizer.minimize(loss, global_step=tf.compat.v1.train.get_global_step())
-        res = tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
+        res = tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op, eval_metric_ops=metrics)
         return res