text: # coding=utf-8
  # Copyright 2020 The Uncertainty Baselines Authors.
  #
  # 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.
  
  """Wide ResNet 28-10 on CIFAR-10/100 trained with maximum likelihood.
  
  Hyperparameters differ slightly from the original paper's code
  (https://github.com/szagoruyko/wide-residual-networks) as TensorFlow uses, for
  example, l2 instead of weight decay, and a different parameterization for SGD's
  momentum.
  """
  
  import functools
  import os
  import time
  from absl import app
  from absl import flags
  from absl import logging
  
  import tensorflow as tf
  import tensorflow_datasets as tfds
  import uncertainty_baselines as ub
  import utils  # local file import
  import uncertainty_metrics as um
  
  flags.DEFINE_integer('seed', 42, 'Random seed.')
  flags.DEFINE_integer('per_core_batch_size', 64, 'Batch size per TPU core/GPU.')
  flags.DEFINE_float('base_learning_rate', 0.1,
                     'Base learning rate when total batch size is 128. It is '
                     'scaled by the ratio of the total batch size to 128.')
  flags.DEFINE_integer('lr_warmup_epochs', 1,
                       'Number of epochs for a linear warmup to the initial '
                       'learning rate. Use 0 to do no warmup.')
  flags.DEFINE_float('lr_decay_ratio', 0.2, 'Amount to decay learning rate.')
  flags.DEFINE_list('lr_decay_epochs', ['60', '120', '160'],
                    'Epochs to decay learning rate by.')
  flags.DEFINE_float('l2', 2e-4, 'L2 regularization coefficient.')
  flags.DEFINE_enum('dataset', 'cifar10',
                    enum_values=['cifar10', 'cifar100'],
                    help='Dataset.')
  # TODO(ghassen): consider adding CIFAR-100-C to TFDS.
  flags.DEFINE_string('cifar100_c_path', None,
                      'Path to the TFRecords files for CIFAR-100-C. Only valid '
                      '(and required) if dataset is cifar100 and corruptions.')
  flags.DEFINE_integer('corruptions_interval', 200,
                       'Number of epochs between evaluating on the corrupted '
                       'test data. Use -1 to never evaluate.')
  flags.DEFINE_integer('checkpoint_interval', 25,
                       'Number of epochs between saving checkpoints. Use -1 to '
                       'never save checkpoints.')
  flags.DEFINE_integer('num_bins', 15, 'Number of bins for ECE.')
  flags.DEFINE_string(
      'data_dir',
      None,
      'Directory where the dataset is stored to be loaded via tfds.load. '
      'Optional, useful for loading datasets stored on GCS.')
  flags.DEFINE_string('output_dir', '/tmp/cifar', 'Output directory.')
  flags.DEFINE_bool(
      'common_output_dir',
      False,
      'Whether the output directory has the work unit id appended to it or not.')
  flags.DEFINE_integer('train_epochs', 200, 'Number of training epochs.')
  
  # Accelerator flags.
  flags.DEFINE_bool('use_gpu', False, 'Whether to run on GPU or otherwise TPU.')
  flags.DEFINE_bool('use_bfloat16', False, 'Whether to use mixed precision.')
  flags.DEFINE_integer('num_cores', 8, 'Number of TPU cores or number of GPUs.')
  flags.DEFINE_string('tpu', None,
                      'Name of the TPU. Only used if use_gpu is False.')
  FLAGS = flags.FLAGS
  
  
  def main(argv):
    del argv  # unused arg
    tf.io.gfile.makedirs(FLAGS.output_dir)
    logging.info('Saving checkpoints at %s', FLAGS.output_dir)
    tf.random.set_seed(FLAGS.seed)
  
    if FLAGS.use_gpu:
      logging.info('Use GPU')
      strategy = tf.distribute.MirroredStrategy()
    else:
      logging.info('Use TPU at %s',
                   FLAGS.tpu if FLAGS.tpu is not None else 'local')
      resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu)
      tf.config.experimental_connect_to_cluster(resolver)
      tf.tpu.experimental.initialize_tpu_system(resolver)
      strategy = tf.distribute.experimental.TPUStrategy(resolver)
  
    train_input_fn = utils.load_input_fn(
        split=tfds.Split.TRAIN,
        name=FLAGS.dataset,
        batch_size=FLAGS.per_core_batch_size,
        use_bfloat16=FLAGS.use_bfloat16,
        data_dir=FLAGS.data_dir)
    clean_test_input_fn = utils.load_input_fn(
        split=tfds.Split.TEST,
        name=FLAGS.dataset,
        batch_size=FLAGS.per_core_batch_size,
        use_bfloat16=FLAGS.use_bfloat16,
        data_dir=FLAGS.data_dir)
    train_dataset = strategy.experimental_distribute_datasets_from_function(
        train_input_fn)
    test_datasets = {
        'clean': strategy.experimental_distribute_datasets_from_function(
            clean_test_input_fn),
    }
    if FLAGS.corruptions_interval > 0:
      if FLAGS.dataset == 'cifar10':
        load_c_input_fn = utils.load_cifar10_c_input_fn
      else:
        load_c_input_fn = functools.partial(utils.load_cifar100_c_input_fn,
                                            path=FLAGS.cifar100_c_path)
      corruption_types, max_intensity = utils.load_corrupted_test_info(
          FLAGS.dataset)
      for corruption in corruption_types:
        for intensity in range(1, max_intensity + 1):
          input_fn = load_c_input_fn(
              corruption_name=corruption,
              corruption_intensity=intensity,
              batch_size=FLAGS.per_core_batch_size,
              use_bfloat16=FLAGS.use_bfloat16)
          test_datasets['{0}_{1}'.format(corruption, intensity)] = (
              strategy.experimental_distribute_datasets_from_function(input_fn))
  
    ds_info = tfds.builder(FLAGS.dataset).info
    batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores
    steps_per_epoch = ds_info.splits['train'].num_examples // batch_size
    steps_per_eval = ds_info.splits['test'].num_examples // batch_size
    num_classes = ds_info.features['label'].num_classes
  
    if FLAGS.use_bfloat16:
      policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16')
      tf.keras.mixed_precision.experimental.set_policy(policy)
  
    summary_writer = tf.summary.create_file_writer(
        os.path.join(FLAGS.output_dir, 'summaries'))
  
    with strategy.scope():
      logging.info('Building ResNet model')
      model = ub.models.wide_resnet(
          input_shape=ds_info.features['image'].shape,
          depth=28,
          width_multiplier=10,
          num_classes=num_classes,
          l2=FLAGS.l2,
          version=2)
      logging.info('Model input shape: %s', model.input_shape)
      logging.info('Model output shape: %s', model.output_shape)
      logging.info('Model number of weights: %s', model.count_params())
      # Linearly scale learning rate and the decay epochs by vanilla settings.
      base_lr = FLAGS.base_learning_rate * batch_size / 128
      lr_decay_epochs = [(int(start_epoch_str) * FLAGS.train_epochs) // 200
                         for start_epoch_str in FLAGS.lr_decay_epochs]
      lr_schedule = utils.LearningRateSchedule(
          steps_per_epoch,
          base_lr,
          decay_ratio=FLAGS.lr_decay_ratio,
          decay_epochs=lr_decay_epochs,
          warmup_epochs=FLAGS.lr_warmup_epochs)
      optimizer = tf.keras.optimizers.SGD(lr_schedule,
                                          momentum=0.9,
                                          nesterov=True)
      metrics = {
          'train/negative_log_likelihood':
              tf.keras.metrics.Mean(),
          'train/accuracy':
              tf.keras.metrics.SparseCategoricalAccuracy(),
          'train/loss':
              tf.keras.metrics.Mean(),
          'train/ece':
              um.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
          'test/negative_log_likelihood':
              tf.keras.metrics.Mean(),
          'test/accuracy':
              tf.keras.metrics.SparseCategoricalAccuracy(),
          'test/ece':
              um.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
      }
      if FLAGS.corruptions_interval > 0:
        corrupt_metrics = {}
        for intensity in range(1, max_intensity + 1):
          for corruption in corruption_types:
            dataset_name = '{0}_{1}'.format(corruption, intensity)
            corrupt_metrics['test/nll_{}'.format(dataset_name)] = (
                tf.keras.metrics.Mean())
            corrupt_metrics['test/accuracy_{}'.format(dataset_name)] = (
                tf.keras.metrics.SparseCategoricalAccuracy())
            corrupt_metrics['test/ece_{}'.format(dataset_name)] = (
                um.ExpectedCalibrationError(num_bins=FLAGS.num_bins))
  
      checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
      latest_checkpoint = tf.train.latest_checkpoint(FLAGS.output_dir)
      initial_epoch = 0
      if latest_checkpoint:
        # checkpoint.restore must be within a strategy.scope() so that optimizer
        # slot variables are mirrored.
        checkpoint.restore(latest_checkpoint)
        logging.info('Loaded checkpoint %s', latest_checkpoint)
        initial_epoch = optimizer.iterations.numpy() // steps_per_epoch
  
    @tf.function
    def train_step(iterator):
      """Training StepFn."""
      def step_fn(inputs):
        """Per-Replica StepFn."""
        images, labels = inputs
        with tf.GradientTape() as tape:
          logits = model(images, training=True)
          if FLAGS.use_bfloat16:
            logits = tf.cast(logits, tf.float32)
          negative_log_likelihood = tf.reduce_mean(
              tf.keras.losses.sparse_categorical_crossentropy(labels,
                                                              logits,
                                                              from_logits=True))
          l2_loss = sum(model.losses)
          loss = negative_log_likelihood + l2_loss
          # Scale the loss given the TPUStrategy will reduce sum all gradients.
          scaled_loss = loss / strategy.num_replicas_in_sync
  
        grads = tape.gradient(scaled_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
  
        probs = tf.nn.softmax(logits)
        metrics['train/ece'].update_state(labels, probs)
        metrics['train/loss'].update_state(loss)
        metrics['train/negative_log_likelihood'].update_state(
            negative_log_likelihood)
        metrics['train/accuracy'].update_state(labels, logits)
  
      strategy.run(step_fn, args=(next(iterator),))
  
    @tf.function
    def test_step(iterator, dataset_name):
      """Evaluation StepFn."""
      def step_fn(inputs):
        """Per-Replica StepFn."""
        images, labels = inputs
        logits = model(images, training=False)
        if FLAGS.use_bfloat16:
          logits = tf.cast(logits, tf.float32)
        probs = tf.nn.softmax(logits)
        negative_log_likelihood = tf.reduce_mean(
            tf.keras.losses.sparse_categorical_crossentropy(labels, probs))
  
        if dataset_name == 'clean':
          metrics['test/negative_log_likelihood'].update_state(
              negative_log_likelihood)
          metrics['test/accuracy'].update_state(labels, probs)
          metrics['test/ece'].update_state(labels, probs)
        else:
          corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state(
              negative_log_likelihood)
          corrupt_metrics['test/accuracy_{}'.format(dataset_name)].update_state(
              labels, probs)
          corrupt_metrics['test/ece_{}'.format(dataset_name)].update_state(
              labels, probs)
  
      strategy.run(step_fn, args=(next(iterator),))
  
    metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()})
  
    train_iterator = iter(train_dataset)
    start_time = time.time()
    for epoch in range(initial_epoch, FLAGS.train_epochs):
      logging.info('Starting to run epoch: %s', epoch)
      for step in range(steps_per_epoch):
        train_step(train_iterator)
  
        current_step = epoch * steps_per_epoch + (step + 1)
        max_steps = steps_per_epoch * FLAGS.train_epochs
        time_elapsed = time.time() - start_time
        steps_per_sec = float(current_step) / time_elapsed
        eta_seconds = (max_steps - current_step) / steps_per_sec
        message = ('{:.1%} completion: epoch {:d}/{:d}. {:.1f} steps/s. '
                   'ETA: {:.0f} min. Time elapsed: {:.0f} min'.format(
                       current_step / max_steps,
                       epoch + 1,
                       FLAGS.train_epochs,
                       steps_per_sec,
                       eta_seconds / 60,
                       time_elapsed / 60))
        if step % 20 == 0:
          logging.info(message)
  
      datasets_to_evaluate = {'clean': test_datasets['clean']}
      if (FLAGS.corruptions_interval > 0 and
          (epoch + 1) % FLAGS.corruptions_interval == 0):
        datasets_to_evaluate = test_datasets
      for dataset_name, test_dataset in datasets_to_evaluate.items():
        test_iterator = iter(test_dataset)
        logging.info('Testing on dataset %s', dataset_name)
        for step in range(steps_per_eval):
          if step % 20 == 0:
            logging.info('Starting to run eval step %s of epoch: %s', step,
                         epoch)
          test_start_time = time.time()
          test_step(test_iterator, dataset_name)
          ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size
          metrics['test/ms_per_example'].update_state(ms_per_example)
  
        logging.info('Done with testing on %s', dataset_name)
  
      corrupt_results = {}
      if (FLAGS.corruptions_interval > 0 and
          (epoch + 1) % FLAGS.corruptions_interval == 0):
        corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics,
                                                          corruption_types,
                                                          max_intensity)
  
      logging.info('Train Loss: %.4f, Accuracy: %.2f%%',
                   metrics['train/loss'].result(),
                   metrics['train/accuracy'].result() * 100)
      logging.info('Test NLL: %.4f, Accuracy: %.2f%%',
                   metrics['test/negative_log_likelihood'].result(),
                   metrics['test/accuracy'].result() * 100)
      total_results = {name: metric.result() for name, metric in metrics.items()}
      total_results.update(corrupt_results)
      with summary_writer.as_default():
        for name, result in total_results.items():
          tf.summary.scalar(name, result, step=epoch + 1)
  
      for metric in metrics.values():
        metric.reset_states()
  
      if (FLAGS.checkpoint_interval > 0 and
          (epoch + 1) % FLAGS.checkpoint_interval == 0):
        checkpoint_name = checkpoint.save(
            os.path.join(FLAGS.output_dir, 'checkpoint'))
        logging.info('Saved checkpoint to %s', checkpoint_name)
  
    final_checkpoint_name = checkpoint.save(
        os.path.join(FLAGS.output_dir, 'checkpoint'))
    logging.info('Saved last checkpoint to %s', final_checkpoint_name)
  
  if __name__ == '__main__':
    app.run(main)
  