from tensorflow.contrib.learn.python.learn.summary_writer_cache import SummaryWriterCache


  def __init__(self, every_n_steps=100, output_dir=None,
               summary_writer=None):
    super(StepCounter, self).__init__(every_n_steps=every_n_steps)
    self._summary_tag = "global_step/sec"
    self._last_reported_step = None
    self._last_reported_time = None
    self._summary_writer = summary_writer
    if summary_writer is None and output_dir:
      self._summary_writer = SummaryWriterCache.get(output_dir)

  def set_estimator(self, estimator):
    super(StepCounter, self).set_estimator(estimator)
    if self._summary_writer is None:
      self._summary_writer = SummaryWriterCache.get(estimator.model_dir)

  def every_n_step_end(self, current_step, outputs):
    current_time = time.time()
