message: Clean up the TODO for training_context and use on_batch().
diff: @@ -115,14 +115,8 @@ def run_one_epoch(model,
  current_batch_size = batch_size
  else:
  current_batch_size = num_samples - step * batch_size
  -
  - # TODO(scottzhu): Maybe update the training context to take into account
  - # whether a batch of training happens. Then it could still use a
  - # context manager
  - batch_logs = {'batch': step, 'size': current_batch_size}
  - training_context.callbacks._call_batch_hook(
  - mode, 'begin', step, batch_logs)
  - training_context.progbar.on_batch_begin(step, batch_logs)
  + with training_context.on_batch(
  + step=step, mode=mode, size=current_batch_size) as batch_logs:
  try:
  batch_outs = execution_function(iterator)
  except (StopIteration, errors.OutOfRangeError):
  @@ -149,6 +143,9 @@ def run_one_epoch(model,
  'when building your dataset.'.format(
  total_epochs * steps_per_epoch))
  # In either case, break out the loop for training batch.
  + # Also note the training_context that data inputs are exhausted, so all
  + # the post batch hooks can be skipped.
  + batch_logs['data_exhausted'] = True
  break
  if not isinstance(batch_outs, list):
  @@ -168,11 +165,6 @@ def run_one_epoch(model,
  batch_start=step * batch_size,
  batch_end=step * batch_size + current_batch_size)
  cbks.make_logs(model, batch_logs, batch_outs, mode)
  -
  - training_context.callbacks._call_batch_hook(
  - mode, 'end', step, batch_logs)
  - training_context.progbar.on_batch_end(step, batch_logs)
  -
  step += 1
  if callbacks.model.stop_training:
  @@ -180,8 +172,7 @@ def run_one_epoch(model,
  # End of an epoch.
  aggregator.finalize()
  - results = aggregator.results
  - return results
  + return aggregator.results
  class Loop(training_utils.TrainingLoop):
  @@ -623,15 +614,16 @@ class TrainingContext(object):
  self.progbar.on_epoch_end(epoch, epoch_logs)
  @tf_contextlib.contextmanager
  - def on_batch(self, step=0, mode=ModeKeys.TRAIN):
  + def on_batch(self, step=0, mode=ModeKeys.TRAIN, size=1):
  """Provide a scope for running one batch."""
  - batch_logs = {'batch': step, 'size': 1}
  + batch_logs = {'batch': step, 'size': size}
  self.callbacks._call_batch_hook(
  mode, 'begin', step, batch_logs)
  self.progbar.on_batch_begin(step, batch_logs)
  try:
  yield batch_logs
  finally:
  + if not batch_logs.pop('data_exhausted', False):
  self.callbacks._call_batch_hook(
  mode, 'end', step, batch_logs)
  self.progbar.on_batch_end(step, batch_logs)
  