diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 1a34d017eb0..b2fa9eb45ce 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -2480,17 +2480,19 @@ class TPUEstimator(estimator_lib.Estimator): # examples_hook is added to training_hooks for both CPU and TPU # execution. - examples_hook = ExamplesPerSecondHook( - ctx.global_batch_size, - output_dir=self.model_dir, - every_n_steps=self._log_every_n_steps) + if self._log_every_n_steps is not None: + examples_hook = ExamplesPerSecondHook( + ctx.global_batch_size, + output_dir=self.model_dir, + every_n_steps=self._log_every_n_steps) if ctx.is_running_on_cpu(is_export_mode=is_export_mode): logging.info('Running %s on CPU', mode) estimator_spec = model_fn_wrapper.call_without_tpu( features, labels, is_export_mode=is_export_mode) - estimator_spec = estimator_spec._replace( - training_hooks=estimator_spec.training_hooks + (examples_hook,)) + if self._log_every_n_steps is not None: + estimator_spec = estimator_spec._replace( + training_hooks=estimator_spec.training_hooks + (examples_hook,)) return estimator_spec assert labels is None, '`labels` passed to `model_fn` must be `None`.' @@ -2540,10 +2542,6 @@ class TPUEstimator(estimator_lib.Estimator): with ops.control_dependencies([loss]): global_step = array_ops.identity(training.get_global_step()) hooks = input_hooks + shutdown_hooks - logging_hook_frequency = ( # Divide and round up - (self._log_every_n_steps + - self._config.tpu_config.iterations_per_loop - 1) // - self._config.tpu_config.iterations_per_loop) hooks.extend([ TPUInfeedOutfeedSessionHook( ctx, @@ -2553,16 +2551,22 @@ class TPUEstimator(estimator_lib.Estimator): run_infeed_loop_on_coordinator), rendezvous=self._rendezvous[mode], ), - InstallSignalHandlerHook(), - training.LoggingTensorHook({ - 'loss': array_ops.identity(loss), - 'step': global_step, - }, - every_n_iter=logging_hook_frequency) + InstallSignalHandlerHook() ]) - examples_hook._set_steps_per_run( # pylint: disable=protected-access - self._config.tpu_config.iterations_per_loop) - hooks.append(examples_hook) + if self._log_every_n_steps is not None: + logging_hook_frequency = ( # Divide and round up + (self._log_every_n_steps + + self._config.tpu_config.iterations_per_loop - 1) // + self._config.tpu_config.iterations_per_loop) + hooks.append( + training.LoggingTensorHook({ + 'loss': array_ops.identity(loss), + 'step': global_step, + }, + every_n_iter=logging_hook_frequency)) + examples_hook._set_steps_per_run( # pylint: disable=protected-access + self._config.tpu_config.iterations_per_loop) + hooks.append(examples_hook) if training_hooks: hooks.extend(training_hooks)