Cherry-pick classifier-trainer to r2.2.0.

Author: Allen Wang

official/benchmark/keras_imagenet_benchmark.py

@@ -257,7 +257,6 @@ class RuntimeConfig(Config):
 
   Attributes:
     distribution_strategy: e.g. 'mirrored', 'tpu', etc.
-    enable_eager: Whether or not to enable eager mode.
     enable_xla: Whether or not to enable XLA.
     per_gpu_thread_count: thread count per GPU.
     gpu_threads_enabled: Whether or not GPU threads are enabled.
@@ -272,9 +271,12 @@ class RuntimeConfig(Config):
     all_reduce_alg: Defines the algorithm for performing all-reduce.
     num_packs: Sets `num_packs` in the cross device ops used in
       MirroredStrategy.  For details, see tf.distribute.NcclAllReduce.
+    loss_scale: The type of loss scale. This is used when setting the mixed
+      precision policy.
+    run_eagerly: Whether or not to run the experiment eagerly.
+
   """
   distribution_strategy: str = 'mirrored'
-  enable_eager: bool = False
   enable_xla: bool = False
   gpu_threads_enabled: bool = False
   gpu_thread_mode: Optional[str] = None
@@ -286,6 +288,8 @@ class RuntimeConfig(Config):
   task_index: int = -1
   all_reduce_alg: Optional[str] = None
   num_packs: int = 1
+  loss_scale: Optional[str] = None
+  run_eagerly: bool = False
 
 
 @dataclasses.dataclass
@@ -312,7 +316,10 @@ class CallbacksConfig(Config):
       Callback. Defaults to True.
     enable_tensorboard: Whether or not to enable Tensorboard as a Callback.
       Defaults to True.
+    enable_time_history: Whether or not to enable TimeHistory Callbacks.
+      Defaults to True.
 
   """
   enable_checkpoint_and_export: bool = True
   enable_tensorboard: bool = True
+  enable_time_history: bool = True

official/modeling/hyperparams/base_config.py

@@ -1,3 +1,4 @@
+# Lint as: python3
 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,15 +23,20 @@
 from absl import logging
 
 import tensorflow as tf
-from typing import Any, List, MutableMapping, Text
+from typing import Any, List, MutableMapping
+
+from official.utils.misc import keras_utils
 
 
 def get_callbacks(model_checkpoint: bool = True,
                   include_tensorboard: bool = True,
+                  time_history: bool = True,
                   track_lr: bool = True,
                   write_model_weights: bool = True,
                   initial_step: int = 0,
-                  model_dir: Text = None) -> List[tf.keras.callbacks.Callback]:
+                  batch_size: int = 0,
+                  log_steps: int = 0,
+                  model_dir: str = None) -> List[tf.keras.callbacks.Callback]:
   """Get all callbacks."""
   model_dir = model_dir or ''
   callbacks = []
@@ -44,6 +50,11 @@ def get_callbacks(model_checkpoint: bool = True,
         track_lr=track_lr,
         initial_step=initial_step,
         write_images=write_model_weights))
+  if time_history:
+    callbacks.append(keras_utils.TimeHistory(
+        batch_size,
+        log_steps,
+        logdir=model_dir if include_tensorboard else None))
   return callbacks
 
 
@@ -74,7 +85,7 @@ class CustomTensorBoard(tf.keras.callbacks.TensorBoard):
   # classification loss
 
   def __init__(self,
-               log_dir: Text,
+               log_dir: str,
                track_lr: bool = False,
                initial_step: int = 0,
                **kwargs):
@@ -84,7 +95,7 @@ def __init__(self,
 
   def on_batch_begin(self,
                      epoch: int,
-                     logs: MutableMapping[Text, Any] = None) -> None:
+                     logs: MutableMapping[str, Any] = None) -> None:
     self.step += 1
     if logs is None:
       logs = {}
@@ -93,7 +104,7 @@ def on_batch_begin(self,
 
   def on_epoch_begin(self,
                      epoch: int,
-                     logs: MutableMapping[Text, Any] = None) -> None:
+                     logs: MutableMapping[str, Any] = None) -> None:
     if logs is None:
       logs = {}
     metrics = self._calculate_metrics()
@@ -104,14 +115,14 @@ def on_epoch_begin(self,
 
   def on_epoch_end(self,
                    epoch: int,
-                   logs: MutableMapping[Text, Any] = None) -> None:
+                   logs: MutableMapping[str, Any] = None) -> None:
     if logs is None:
       logs = {}
     metrics = self._calculate_metrics()
     logs.update(metrics)
     super(CustomTensorBoard, self).on_epoch_end(epoch, logs)
 
-  def _calculate_metrics(self) -> MutableMapping[Text, Any]:
+  def _calculate_metrics(self) -> MutableMapping[str, Any]:
     logs = {}
     if self._track_lr:
       logs['learning_rate'] = self._calculate_lr()

official/vision/image_classification/callbacks.py

@@ -44,10 +44,24 @@
 from official.vision.image_classification.resnet import common
 from official.vision.image_classification.resnet import resnet_model
 
-MODELS = {
-    'efficientnet': efficientnet_model.EfficientNet.from_name,
-    'resnet': resnet_model.resnet50,
-}
+
+def get_models() -> Mapping[str, tf.keras.Model]:
+  """Returns the mapping from model type name to Keras model."""
+  return  {
+      'efficientnet': efficientnet_model.EfficientNet.from_name,
+      'resnet': resnet_model.resnet50,
+  }
+
+
+def get_dtype_map() -> Mapping[str, tf.dtypes.DType]:
+  """Returns the mapping from dtype string representations to TF dtypes."""
+  return {
+      'float32': tf.float32,
+      'bfloat16': tf.bfloat16,
+      'float16': tf.float16,
+      'fp32': tf.float32,
+      'bf16': tf.bfloat16,
+    }
 
 
 def _get_metrics(one_hot: bool) -> Mapping[Text, Any]:
@@ -120,7 +134,7 @@ def _get_dataset_builders(params: base_configs.ExperimentConfig,
 def get_loss_scale(params: base_configs.ExperimentConfig,
                    fp16_default: float = 128.) -> float:
   """Returns the loss scale for initializations."""
-  loss_scale = params.model.loss.loss_scale
+  loss_scale = params.runtime.loss_scale
   if loss_scale == 'dynamic':
     return loss_scale
   elif loss_scale is not None:
@@ -145,7 +159,7 @@ def _get_params_from_flags(flags_obj: flags.FlagValues):
           'name': model,
       },
       'runtime': {
-          'enable_eager': flags_obj.enable_eager,
+          'run_eagerly': flags_obj.run_eagerly,
           'tpu': flags_obj.tpu,
       },
       'train_dataset': {
@@ -154,8 +168,10 @@ def _get_params_from_flags(flags_obj: flags.FlagValues):
       'validation_dataset': {
           'data_dir': flags_obj.data_dir,
       },
-      'test_dataset': {
-          'data_dir': flags_obj.data_dir,
+      'train': {
+          'time_history': {
+              'log_steps': flags_obj.log_steps,
+          },
       },
   }
 
@@ -212,10 +228,11 @@ def resume_from_checkpoint(model: tf.keras.Model,
   return int(initial_epoch)
 
 
-def initialize(params: base_configs.ExperimentConfig):
+def initialize(params: base_configs.ExperimentConfig,
+               dataset_builder: dataset_factory.DatasetBuilder):
   """Initializes backend related initializations."""
   keras_utils.set_session_config(
-      enable_eager=params.runtime.enable_eager,
+      enable_eager=params.runtime.run_eagerly,
       enable_xla=params.runtime.enable_xla)
   if params.runtime.gpu_threads_enabled:
     keras_utils.set_gpu_thread_mode_and_count(
@@ -224,20 +241,19 @@ def initialize(params: base_configs.ExperimentConfig):
         num_gpus=params.runtime.num_gpus,
         datasets_num_private_threads=params.runtime.dataset_num_private_threads)
 
-  dataset = params.train_dataset or params.validation_dataset
-  performance.set_mixed_precision_policy(dataset.dtype)
+  performance.set_mixed_precision_policy(dataset_builder.dtype)
 
-  if dataset.data_format:
-    data_format = dataset.data_format
-  elif tf.config.list_physical_devices('GPU'):
+  if dataset_builder.config.data_format:
+    data_format = dataset_builder.config.data_format
+  if tf.config.list_physical_devices('GPU'):
     data_format = 'channels_first'
   else:
     data_format = 'channels_last'
   tf.keras.backend.set_image_data_format(data_format)
   distribution_utils.configure_cluster(
       params.runtime.worker_hosts,
       params.runtime.task_index)
-  if params.runtime.enable_eager:
+  if params.runtime.run_eagerly:
     # Enable eager execution to allow step-by-step debugging
     tf.config.experimental_run_functions_eagerly(True)
 
@@ -254,7 +270,7 @@ def define_classifier_flags():
       default=None,
       help='Mode to run: `train`, `eval`, `train_and_eval` or `export`.')
   flags.DEFINE_bool(
-      'enable_eager',
+      'run_eagerly',
       default=None,
       help='Use eager execution and disable autograph for debugging.')
   flags.DEFINE_string(
@@ -265,6 +281,10 @@ def define_classifier_flags():
       'dataset',
       default=None,
       help='The name of the dataset, e.g. ImageNet, etc.')
+  flags.DEFINE_integer(
+      'log_steps',
+      default=100,
+      help='The interval of steps between logging of batch level stats.')
 
 
 def serialize_config(params: base_configs.ExperimentConfig,
@@ -307,11 +327,13 @@ def train_and_eval(
   train_steps = params.train.steps or train_builder.num_steps
   validation_steps = params.evaluation.steps or validation_builder.num_steps
 
+  initialize(params, train_builder)
+
   logging.info('Global batch size: %d', train_builder.global_batch_size)
 
   with strategy_scope:
     model_params = params.model.model_params.as_dict()
-    model = MODELS[params.model.name](**model_params)
+    model = get_models()[params.model.name](**model_params)
     learning_rate = optimizer_factory.build_learning_rate(
         params=params.model.learning_rate,
         batch_size=train_builder.global_batch_size,
@@ -331,8 +353,7 @@ def train_and_eval(
       loss_obj = tf.keras.losses.SparseCategoricalCrossentropy()
     model.compile(optimizer=optimizer,
                   loss=loss_obj,
-                  metrics=metrics,
-                  run_eagerly=params.runtime.enable_eager)
+                  metrics=metrics)
 
     initial_epoch = 0
     if params.train.resume_checkpoint:
@@ -345,26 +366,37 @@ def train_and_eval(
   callbacks = custom_callbacks.get_callbacks(
       model_checkpoint=params.train.callbacks.enable_checkpoint_and_export,
       include_tensorboard=params.train.callbacks.enable_tensorboard,
+      time_history=params.train.callbacks.enable_time_history,
       track_lr=params.train.tensorboard.track_lr,
       write_model_weights=params.train.tensorboard.write_model_weights,
       initial_step=initial_epoch * train_steps,
+      batch_size=train_builder.global_batch_size,
+      log_steps=params.train.time_history.log_steps,
       model_dir=params.model_dir)
 
+  if params.evaluation.skip_eval:
+    validation_kwargs = {}
+  else:
+    validation_kwargs = {
+        'validation_data': validation_dataset,
+        'validation_steps': validation_steps,
+        'validation_freq': params.evaluation.epochs_between_evals,
+    }
+
   history = model.fit(
       train_dataset,
       epochs=train_epochs,
       steps_per_epoch=train_steps,
       initial_epoch=initial_epoch,
       callbacks=callbacks,
-      validation_data=validation_dataset,
-      validation_steps=validation_steps,
-      validation_freq=params.evaluation.epochs_between_evals)
+      **validation_kwargs)
 
-  validation_output = model.evaluate(
-      validation_dataset, steps=validation_steps, verbose=2)
+  validation_output = None
+  if not params.evaluation.skip_eval:
+    validation_output = model.evaluate(
+        validation_dataset, steps=validation_steps, verbose=2)
 
   # TODO(dankondratyuk): eval and save final test accuracy
-
   stats = common.build_stats(history,
                              validation_output,
                              callbacks)
@@ -375,7 +407,7 @@ def export(params: base_configs.ExperimentConfig):
   """Runs the model export functionality."""
   logging.info('Exporting model.')
   model_params = params.model.model_params.as_dict()
-  model = MODELS[params.model.name](**model_params)
+  model = get_models()[params.model.name](**model_params)
   checkpoint = params.export.checkpoint
   if checkpoint is None:
     logging.info('No export checkpoint was provided. Using the latest '
@@ -398,8 +430,6 @@ def run(flags_obj: flags.FlagValues,
     Dictionary of training/eval stats
   """
   params = _get_params_from_flags(flags_obj)
-  initialize(params)
-
   if params.mode == 'train_and_eval':
     return train_and_eval(params, strategy_override)
   elif params.mode == 'export_only':

official/vision/image_classification/classifier_trainer.py

@@ -239,8 +239,8 @@ def test_get_model_size(self, model, model_name, expected):
   )
   def test_get_loss_scale(self, loss_scale, dtype, expected):
     config = base_configs.ExperimentConfig(
-        model=base_configs.ModelConfig(
-            loss=base_configs.LossConfig(loss_scale=loss_scale)),
+        runtime=base_configs.RuntimeConfig(
+            loss_scale=loss_scale),
         train_dataset=dataset_factory.DatasetConfig(dtype=dtype))
     ls = classifier_trainer.get_loss_scale(config, fp16_default=128)
     self.assertEqual(ls, expected)
@@ -252,7 +252,7 @@ def test_get_loss_scale(self, loss_scale, dtype, expected):
   def test_initialize(self, dtype):
     config = base_configs.ExperimentConfig(
         runtime=base_configs.RuntimeConfig(
-            enable_eager=False,
+            run_eagerly=False,
             enable_xla=False,
             gpu_threads_enabled=True,
             per_gpu_thread_count=1,
@@ -264,7 +264,14 @@ def test_initialize(self, dtype):
         model=base_configs.ModelConfig(
             loss=base_configs.LossConfig(loss_scale='dynamic')),
     )
-    classifier_trainer.initialize(config)
+
+    class EmptyClass:
+      pass
+    fake_ds_builder = EmptyClass()
+    fake_ds_builder.dtype = dtype
+    fake_ds_builder.config = EmptyClass()
+    fake_ds_builder.config.data_format = None
+    classifier_trainer.initialize(config, fake_ds_builder)
 
   def test_resume_from_checkpoint(self):
     """Tests functionality for resuming from checkpoint."""