Simplified estimator logic by MonitoredSession.

Removed graph_action usage.
Change: 144126485

Author: ispirmustafa

tensorflow/contrib/factorization/python/ops/gmm.py

@@ -21,20 +21,28 @@
 from __future__ import division
 from __future__ import print_function
 
+import time
+
 import numpy as np
 
 from tensorflow.contrib import framework
 from tensorflow.contrib.factorization.python.ops import gmm_ops
 from tensorflow.contrib.framework.python.framework import checkpoint_utils
 from tensorflow.contrib.framework.python.ops import variables
-from tensorflow.contrib.learn.python.learn.estimators import estimator
+from tensorflow.contrib.learn.python.learn import graph_actions
+from tensorflow.contrib.learn.python.learn import monitors as monitor_lib
+from tensorflow.contrib.learn.python.learn.estimators import estimator as estimator_lib
+from tensorflow.contrib.learn.python.learn.estimators import model_fn as model_fn_lib
 from tensorflow.contrib.learn.python.learn.estimators._sklearn import TransformerMixin
 from tensorflow.contrib.learn.python.learn.learn_io import data_feeder
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed as random_seed_lib
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops.control_flow_ops import with_dependencies
+from tensorflow.python.platform import tf_logging as logging
 
 
 def _streaming_sum(scalar_tensor):
@@ -44,7 +52,7 @@ def _streaming_sum(scalar_tensor):
   return sum_metric, sum_update
 
 
-class GMM(estimator.Estimator, TransformerMixin):
+class GMM(estimator_lib.Estimator, TransformerMixin):
   """GMM clustering."""
   SCORES = 'scores'
   ASSIGNMENTS = 'assignments'
@@ -116,7 +124,8 @@ def fit(self, x, y=None, monitors=None, logdir=None, steps=None):
     self._data_feeder = data_feeder.setup_train_data_feeder(x, None,
                                                             self._num_clusters,
                                                             self.batch_size)
-    self._train_model(
+    _legacy_train_model(  # pylint: disable=protected-access
+        self,
         input_fn=self._data_feeder.input_builder,
         feed_fn=self._data_feeder.get_feed_dict_fn(),
         steps=steps or self.steps,
@@ -218,3 +227,90 @@ def _get_eval_ops(self, features, _, unused_metrics):
          self._covariance_type,
          self._params)
     return {GMM.SCORES: _streaming_sum(math_ops.reduce_sum(losses))}
+
+
+# TODO(xavigonzalvo): delete this after implementing model-fn based Estimator.
+def _legacy_train_model(estimator,
+                        input_fn,
+                        steps,
+                        feed_fn=None,
+                        init_op=None,
+                        init_feed_fn=None,
+                        init_fn=None,
+                        device_fn=None,
+                        monitors=None,
+                        log_every_steps=100,
+                        fail_on_nan_loss=True,
+                        max_steps=None):
+  """Legacy train function of Estimator."""
+  if hasattr(estimator.config, 'execution_mode'):
+    if estimator.config.execution_mode not in ('all', 'train'):
+      return
+
+    # Stagger startup of worker sessions based on task id.
+    sleep_secs = min(
+        estimator.config.training_worker_max_startup_secs,
+        estimator.config.task_id *
+        estimator.config.training_worker_session_startup_stagger_secs)
+    if sleep_secs:
+      logging.info('Waiting %d secs before starting task %d.', sleep_secs,
+                   estimator.config.task_id)
+      time.sleep(sleep_secs)
+
+  # Device allocation
+  device_fn = device_fn or estimator._device_fn  # pylint: disable=protected-access
+
+  with ops.Graph().as_default() as g, g.device(device_fn):
+    random_seed_lib.set_random_seed(estimator.config.tf_random_seed)
+    global_step = framework.create_global_step(g)
+    features, labels = input_fn()
+    estimator._check_inputs(features, labels)  # pylint: disable=protected-access
+
+    # The default return type of _get_train_ops is ModelFnOps. But there are
+    # some subclasses of tf.contrib.learn.Estimator which override this
+    # method and use the legacy signature, namely _get_train_ops returns a
+    # (train_op, loss) tuple. The following else-statement code covers these
+    # cases, but will soon be deleted after the subclasses are updated.
+    # TODO(b/32664904): Update subclasses and delete the else-statement.
+    train_ops = estimator._get_train_ops(features, labels)  # pylint: disable=protected-access
+    if isinstance(train_ops, model_fn_lib.ModelFnOps):  # Default signature
+      train_op = train_ops.train_op
+      loss_op = train_ops.loss
+      if estimator.config.is_chief:
+        hooks = train_ops.training_chief_hooks + train_ops.training_hooks
+      else:
+        hooks = train_ops.training_hooks
+    else:  # Legacy signature
+      if len(train_ops) != 2:
+        raise ValueError('Expected a tuple of train_op and loss, got {}'.format(
+            train_ops))
+      train_op = train_ops[0]
+      loss_op = train_ops[1]
+      hooks = []
+
+    hooks += monitor_lib.replace_monitors_with_hooks(monitors, estimator)
+
+    ops.add_to_collection(ops.GraphKeys.LOSSES, loss_op)
+    return graph_actions._monitored_train(  # pylint: disable=protected-access
+        graph=g,
+        output_dir=estimator.model_dir,
+        train_op=train_op,
+        loss_op=loss_op,
+        global_step_tensor=global_step,
+        init_op=init_op,
+        init_feed_dict=init_feed_fn() if init_feed_fn is not None else None,
+        init_fn=init_fn,
+        log_every_steps=log_every_steps,
+        supervisor_is_chief=estimator.config.is_chief,
+        supervisor_master=estimator.config.master,
+        supervisor_save_model_secs=estimator.config.save_checkpoints_secs,
+        supervisor_save_model_steps=estimator.config.save_checkpoints_steps,
+        supervisor_save_summaries_steps=estimator.config.save_summary_steps,
+        keep_checkpoint_max=estimator.config.keep_checkpoint_max,
+        keep_checkpoint_every_n_hours=(
+            estimator.config.keep_checkpoint_every_n_hours),
+        feed_fn=feed_fn,
+        steps=steps,
+        fail_on_nan_loss=fail_on_nan_loss,
+        hooks=hooks,
+        max_steps=max_steps)