text: # Lint as: python3
  # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Runs an Image Classification model."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  
  import pprint
  from typing import Any, Tuple, Text, Optional, Mapping
  
  from absl import app
  from absl import flags
  from absl import logging
  import tensorflow as tf
  
  from official.modeling import performance
  from official.modeling.hyperparams import params_dict
  from official.utils import hyperparams_flags
  from official.utils.logs import logger
  from official.utils.misc import distribution_utils
  from official.utils.misc import keras_utils
  from official.vision.image_classification import callbacks as custom_callbacks
  from official.vision.image_classification import dataset_factory
  from official.vision.image_classification import optimizer_factory
  from official.vision.image_classification.configs import base_configs
  from official.vision.image_classification.configs import configs
  from official.vision.image_classification.efficientnet import efficientnet_model
  from official.vision.image_classification.resnet import common
  from official.vision.image_classification.resnet import resnet_model
  
  
  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]:
    """Get a dict of available metrics to track."""
    if one_hot:
      return {
          # (name, metric_fn)
          'acc': tf.keras.metrics.CategoricalAccuracy(name='accuracy'),
          'accuracy': tf.keras.metrics.CategoricalAccuracy(name='accuracy'),
          'top_1': tf.keras.metrics.CategoricalAccuracy(name='accuracy'),
          'top_5': tf.keras.metrics.TopKCategoricalAccuracy(
              k=5,
              name='top_5_accuracy'),
      }
    else:
      return {
          # (name, metric_fn)
          'acc': tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
          'accuracy': tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
          'top_1': tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
          'top_5': tf.keras.metrics.SparseTopKCategoricalAccuracy(
              k=5,
              name='top_5_accuracy'),
      }
  
  
  def get_image_size_from_model(
      params: base_configs.ExperimentConfig) -> Optional[int]:
    """If the given model has a preferred image size, return it."""
    if params.model_name == 'efficientnet':
      efficientnet_name = params.model.model_params.model_name
      if efficientnet_name in efficientnet_model.MODEL_CONFIGS:
        return efficientnet_model.MODEL_CONFIGS[efficientnet_name].resolution
    return None
  
  
  def _get_dataset_builders(params: base_configs.ExperimentConfig,
                            strategy: tf.distribute.Strategy,
                            one_hot: bool
                           ) -> Tuple[Any, Any, Any]:
    """Create and return train, validation, and test dataset builders."""
    if one_hot:
      logging.warning('label_smoothing > 0, so datasets will be one hot encoded.')
    else:
      logging.warning('label_smoothing not applied, so datasets will not be one '
                      'hot encoded.')
  
    num_devices = strategy.num_replicas_in_sync
    image_size = get_image_size_from_model(params)
  
    dataset_configs = [
        params.train_dataset, params.validation_dataset
    ]
    builders = []
  
    for config in dataset_configs:
      if config is not None and config.has_data:
        builder = dataset_factory.DatasetBuilder(
            config,
            image_size=image_size or config.image_size,
            num_devices=num_devices,
            one_hot=one_hot)
      else:
        builder = None
      builders.append(builder)
  
    return builders
  
  
  def get_loss_scale(params: base_configs.ExperimentConfig,
                     fp16_default: float = 128.) -> float:
    """Returns the loss scale for initializations."""
    loss_scale = params.runtime.loss_scale
    if loss_scale == 'dynamic':
      return loss_scale
    elif loss_scale is not None:
      return float(loss_scale)
    elif (params.train_dataset.dtype == 'float32' or
          params.train_dataset.dtype == 'bfloat16'):
      return 1.
    else:
      assert params.train_dataset.dtype == 'float16'
      return fp16_default
  
  
  def _get_params_from_flags(flags_obj: flags.FlagValues):
    """Get ParamsDict from flags."""
    model = flags_obj.model_type.lower()
    dataset = flags_obj.dataset.lower()
    params = configs.get_config(model=model, dataset=dataset)
  
    flags_overrides = {
        'model_dir': flags_obj.model_dir,
        'mode': flags_obj.mode,
        'model': {
            'name': model,
        },
        'runtime': {
            'run_eagerly': flags_obj.run_eagerly,
            'tpu': flags_obj.tpu,
        },
        'train_dataset': {
            'data_dir': flags_obj.data_dir,
        },
        'validation_dataset': {
            'data_dir': flags_obj.data_dir,
        },
        'train': {
            'time_history': {
                'log_steps': flags_obj.log_steps,
            },
        },
    }
  
    overriding_configs = (flags_obj.config_file,
                          flags_obj.params_override,
                          flags_overrides)
  
    pp = pprint.PrettyPrinter()
  
    logging.info('Base params: %s', pp.pformat(params.as_dict()))
  
    for param in overriding_configs:
      logging.info('Overriding params: %s', param)
      # Set is_strict to false because we can have dynamic dict parameters.
      params = params_dict.override_params_dict(params, param, is_strict=False)
  
    params.validate()
    params.lock()
  
    logging.info('Final model parameters: %s', pp.pformat(params.as_dict()))
    return params
  
  
  def resume_from_checkpoint(model: tf.keras.Model,
                             model_dir: str,
                             train_steps: int) -> int:
    """Resumes from the latest checkpoint, if possible.
  
    Loads the model weights and optimizer settings from a checkpoint.
    This function should be used in case of preemption recovery.
  
    Args:
      model: The model whose weights should be restored.
      model_dir: The directory where model weights were saved.
      train_steps: The number of steps to train.
  
    Returns:
      The epoch of the latest checkpoint, or 0 if not restoring.
  
    """
    logging.info('Load from checkpoint is enabled.')
    latest_checkpoint = tf.train.latest_checkpoint(model_dir)
    logging.info('latest_checkpoint: %s', latest_checkpoint)
    if not latest_checkpoint:
      logging.info('No checkpoint detected.')
      return 0
  
    logging.info('Checkpoint file %s found and restoring from '
                 'checkpoint', latest_checkpoint)
    model.load_weights(latest_checkpoint)
    initial_epoch = model.optimizer.iterations // train_steps
    logging.info('Completed loading from checkpoint.')
    logging.info('Resuming from epoch %d', initial_epoch)
    return int(initial_epoch)
  
  
  def initialize(params: base_configs.ExperimentConfig,
                 dataset_builder: dataset_factory.DatasetBuilder):
    """Initializes backend related initializations."""
    keras_utils.set_session_config(
        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(
          per_gpu_thread_count=params.runtime.per_gpu_thread_count,
          gpu_thread_mode=params.runtime.gpu_thread_mode,
          num_gpus=params.runtime.num_gpus,
          datasets_num_private_threads=params.runtime.dataset_num_private_threads)
  
    performance.set_mixed_precision_policy(dataset_builder.dtype,
                                           get_loss_scale(params))
    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.run_eagerly:
      # Enable eager execution to allow step-by-step debugging
      tf.config.experimental_run_functions_eagerly(True)
  
  
  def define_classifier_flags():
    """Defines common flags for image classification."""
    hyperparams_flags.initialize_common_flags()
    flags.DEFINE_string(
        'data_dir',
        default=None,
        help='The location of the input data.')
    flags.DEFINE_string(
        'mode',
        default=None,
        help='Mode to run: `train`, `eval`, `train_and_eval` or `export`.')
    flags.DEFINE_bool(
        'run_eagerly',
        default=None,
        help='Use eager execution and disable autograph for debugging.')
    flags.DEFINE_string(
        'model_type',
        default=None,
        help='The type of the model, e.g. EfficientNet, etc.')
    flags.DEFINE_string(
        '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,
                       model_dir: str):
    """Serializes and saves the experiment config."""
    params_save_path = os.path.join(model_dir, 'params.yaml')
    logging.info('Saving experiment configuration to %s', params_save_path)
    tf.io.gfile.makedirs(model_dir)
    params_dict.save_params_dict_to_yaml(params, params_save_path)
  
  
  def train_and_eval(
      params: base_configs.ExperimentConfig,
      strategy_override: tf.distribute.Strategy) -> Mapping[str, Any]:
    """Runs the train and eval path using compile/fit."""
    logging.info('Running train and eval.')
  
    # Note: for TPUs, strategy and scope should be created before the dataset
    strategy = strategy_override or distribution_utils.get_distribution_strategy(
        distribution_strategy=params.runtime.distribution_strategy,
        all_reduce_alg=params.runtime.all_reduce_alg,
        num_gpus=params.runtime.num_gpus,
        tpu_address=params.runtime.tpu)
  
    strategy_scope = distribution_utils.get_strategy_scope(strategy)
  
    logging.info('Detected %d devices.', strategy.num_replicas_in_sync)
  
    label_smoothing = params.model.loss.label_smoothing
    one_hot = label_smoothing and label_smoothing > 0
  
    builders = _get_dataset_builders(params, strategy, one_hot)
    datasets = [builder.build() if builder else None for builder in builders]
  
    # Unpack datasets and builders based on train/val/test splits
    train_builder, validation_builder = builders  # pylint: disable=unbalanced-tuple-unpacking
    train_dataset, validation_dataset = datasets
  
    train_epochs = params.train.epochs
    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 = 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,
          train_steps=train_steps)
      optimizer = optimizer_factory.build_optimizer(
          optimizer_name=params.model.optimizer.name,
          base_learning_rate=learning_rate,
          params=params.model.optimizer.as_dict())
  
      metrics_map = _get_metrics(one_hot)
      metrics = [metrics_map[metric] for metric in params.train.metrics]
  
      if one_hot:
        loss_obj = tf.keras.losses.CategoricalCrossentropy(
            label_smoothing=params.model.loss.label_smoothing)
      else:
        loss_obj = tf.keras.losses.SparseCategoricalCrossentropy()
      model.compile(optimizer=optimizer,
                    loss=loss_obj,
                    metrics=metrics)
  
      initial_epoch = 0
      if params.train.resume_checkpoint:
        initial_epoch = resume_from_checkpoint(model=model,
                                               model_dir=params.model_dir,
                                               train_steps=train_steps)
  
      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)
  
    serialize_config(params=params, 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_kwargs,
        experimental_steps_per_execution=params.train.steps_per_loop,
        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)
    return stats
  
  
  def export(params: base_configs.ExperimentConfig):
    """Runs the model export functionality."""
    logging.info('Exporting model.')
    model_params = params.model.model_params.as_dict()
    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 '
                   'checkpoint from model_dir.')
      checkpoint = tf.train.latest_checkpoint(params.model_dir)
  
    model.load_weights(checkpoint)
    model.save(params.export.destination)
  
  
  def run(flags_obj: flags.FlagValues,
          strategy_override: tf.distribute.Strategy = None) -> Mapping[str, Any]:
    """Runs Image Classification model using native Keras APIs.
  
    Args:
      flags_obj: An object containing parsed flag values.
      strategy_override: A `tf.distribute.Strategy` object to use for model.
  
    Returns:
      Dictionary of training/eval stats
    """
    params = _get_params_from_flags(flags_obj)
    if params.mode == 'train_and_eval':
      return train_and_eval(params, strategy_override)
    elif params.mode == 'export_only':
      export(params)
    else:
      raise ValueError('{} is not a valid mode.'.format(params.mode))
  
  
  def main(_):
    with logger.benchmark_context(flags.FLAGS):
      stats = run(flags.FLAGS)
    if stats:
      logging.info('Run stats:\n%s', stats)
  
  
  if __name__ == '__main__':
    logging.set_verbosity(logging.INFO)
    define_classifier_flags()
    flags.mark_flag_as_required('data_dir')
    flags.mark_flag_as_required('mode')
    flags.mark_flag_as_required('model_type')
    flags.mark_flag_as_required('dataset')
  
    app.run(main)
  