text: # coding=utf-8
  # Copyright 2021 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.
  
  """Variational inference for Wide ResNet 28-10 on CIFAR-10 and CIFAR-100.
  
  This script performs variational inference with a few notable techniques:
  
  1. Normal prior whose mean is tied at the variational posterior's. This makes
     the KL penalty only penalize the weight posterior's standard deviation and
     not its mean. The prior's standard deviation is a fixed hyperparameter.
  2. Fully factorized normal variational distribution (Blundell et al., 2015).
  3. Flipout for lower-variance gradients in convolutional layers and the final
     dense layer (Wen et al., 2018).
  4. KL annealing (Bowman et al., 2015).
  """
  
  import os
  import time
  from absl import app
  from absl import flags
  from absl import logging
  import robustness_metrics as rm
  import tensorflow as tf
  import tensorflow_datasets as tfds
  import uncertainty_baselines as ub
  import utils  # local file import
  from tensorboard.plugins.hparams import api as hp
  from entropy_uncertainty import compute_entropy_uncertainty
  import threading
  
  flags.DEFINE_integer('kl_annealing_epochs', 200,
                       'Number of epoch over which to anneal the KL term to 1.')
  flags.DEFINE_float('prior_stddev', 0.1, 'Fixed stddev for weight prior.')
  flags.DEFINE_float('stddev_init', 1e-3,
                     'Initialization of posterior standard deviation parameters.')
  flags.DEFINE_integer('num_eval_samples', 5,
                       'Number of samples per example during evaluation. Only '
                       'used for corrupted predicitons.')
  # Redefining default values
  flags.FLAGS.set_default('l2', 4e-4)
  flags.FLAGS.set_default('train_epochs', 250)
  flags.FLAGS.set_default('train_proportion', 0.95)
  FLAGS = flags.FLAGS
  
  tf.config.run_functions_eagerly(True) 
  
  def main(argv):
    del argv  # unused arg
    # u = Uncertainty()
    tf.io.gfile.makedirs(FLAGS.output_dir)
    logging.info('Saving checkpoints at %s', FLAGS.output_dir)
    tf.random.set_seed(FLAGS.seed)
  
    data_dir = FLAGS.data_dir
    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.TPUStrategy(resolver)
  
    ds_info = tfds.builder(FLAGS.dataset).info
    batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores
    train_dataset_size = ds_info.splits['train'].num_examples
    steps_per_epoch = train_dataset_size // batch_size
    steps_per_eval = ds_info.splits['test'].num_examples // batch_size
    num_classes = ds_info.features['label'].num_classes
  
    train_builder = ub.datasets.get(
        FLAGS.dataset,
        data_dir=data_dir,
        download_data=FLAGS.download_data,
        split=tfds.Split.TRAIN,
        validation_percent=1. - FLAGS.train_proportion)
    train_dataset = train_builder.load(batch_size=batch_size)
    validation_dataset = None
    steps_per_validation = 0
    if FLAGS.train_proportion < 1.0:
      validation_builder = ub.datasets.get(
          FLAGS.dataset,
          data_dir=data_dir,
          download_data=FLAGS.download_data,
          split=tfds.Split.VALIDATION,
          validation_percent=1. - FLAGS.train_proportion)
      validation_dataset = validation_builder.load(batch_size=batch_size)
      validation_dataset = strategy.experimental_distribute_dataset(
          validation_dataset)
      steps_per_validation = validation_builder.num_examples // batch_size
    clean_test_builder = ub.datasets.get(
        FLAGS.dataset,
        data_dir=data_dir,
        download_data=FLAGS.download_data,
        split=tfds.Split.TEST)
    clean_test_dataset = clean_test_builder.load(batch_size=batch_size)
    train_dataset = strategy.experimental_distribute_dataset(train_dataset)
    test_datasets = {
        'clean': strategy.experimental_distribute_dataset(clean_test_dataset),
    }
    steps_per_epoch = train_builder.num_examples // batch_size
    steps_per_eval = clean_test_builder.num_examples // batch_size
    num_classes = 100 if FLAGS.dataset == 'cifar100' else 10
    if FLAGS.corruptions_interval > 0:
      extra_kwargs = {}
      if FLAGS.dataset == 'cifar100':
        data_dir = FLAGS.cifar100_c_path
      corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
      for corruption_type in corruption_types:
        for severity in range(1, 6):
          dataset = ub.datasets.get(
              f'{FLAGS.dataset}_corrupted',
              corruption_type=corruption_type,
              data_dir=data_dir,
              severity=severity,
              split=tfds.Split.TEST,
              **extra_kwargs).load(batch_size=batch_size)
          test_datasets[f'{corruption_type}_{severity}'] = (
              strategy.experimental_distribute_dataset(dataset))
  
    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_variational(
          input_shape=(32, 32, 3),
          depth=28,
          width_multiplier=10,
          num_classes=num_classes,
          prior_stddev=FLAGS.prior_stddev,
          dataset_size=train_dataset_size,
          stddev_init=FLAGS.stddev_init)
      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 = ub.schedules.WarmUpPiecewiseConstantSchedule(
          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=1.0 - FLAGS.one_minus_momentum,
                                          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': rm.metrics.ExpectedCalibrationError(
              num_bins=FLAGS.num_bins),
          'train/kl': tf.keras.metrics.Mean(),
          'train/kl_scale': tf.keras.metrics.Mean(),
          'test/negative_log_likelihood': tf.keras.metrics.Mean(),
          'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
          'test/ece': rm.metrics.ExpectedCalibrationError(
              num_bins=FLAGS.num_bins),
      }
      if validation_dataset:
        metrics.update({
            'validation/negative_log_likelihood': tf.keras.metrics.Mean(),
            'validation/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
            'validation/ece': rm.metrics.ExpectedCalibrationError(
                num_bins=FLAGS.num_bins),
        })
      if FLAGS.corruptions_interval > 0:
        corrupt_metrics = {}
        for intensity in range(1, 6):
          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)] = (
                rm.metrics.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 = inputs['features']
        labels = inputs['labels']
        with tf.GradientTape() as tape:
          logits = model(images, training=True)
          negative_log_likelihood = tf.reduce_mean(
              tf.keras.losses.sparse_categorical_crossentropy(labels,
                                                              logits,
                                                              from_logits=True))
  
          filtered_variables = []
          for var in model.trainable_variables:
            # Apply l2 on the BN parameters and bias terms. This
            # excludes only fast weight approximate posterior/prior parameters,
            # but pay caution to their naming scheme.
            if 'batch_norm' in var.name or 'bias' in var.name:
              filtered_variables.append(tf.reshape(var, (-1,)))
  
          l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss(
              tf.concat(filtered_variables, axis=0))
          kl = sum(model.losses)
          kl_scale = tf.cast(optimizer.iterations + 1, kl.dtype)
          kl_scale /= steps_per_epoch * FLAGS.kl_annealing_epochs
          kl_scale = tf.minimum(1., kl_scale)
          kl_loss = kl_scale * kl
  
          # Scale the loss given the TPUStrategy will reduce sum all gradients.
          loss = negative_log_likelihood + l2_loss + kl_loss
          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'].add_batch(probs, label=labels)
        metrics['train/loss'].update_state(loss)
        metrics['train/negative_log_likelihood'].update_state(
            negative_log_likelihood)
        metrics['train/kl'].update_state(kl)
        metrics['train/kl_scale'].update_state(kl_scale)
        metrics['train/accuracy'].update_state(labels, logits)
  
      for _ in tf.range(tf.cast(steps_per_epoch, tf.int32)):
        strategy.run(step_fn, args=(next(iterator),))
  
    @tf.function
    def test_step(iterator, dataset_split, dataset_name, num_steps, epoch, threads):
      """Evaluation StepFn."""
      def step_fn(inputs):
        """Per-Replica StepFn."""
        images = inputs['features']
        labels = inputs['labels']
        # TODO(trandustin): Use more eval samples only on corrupted predictions;
        # it's expensive but a one-time compute if scheduled post-training.
        # if FLAGS.num_eval_samples > 1 and dataset_name != 'clean':
        if FLAGS.num_eval_samples > 1:
          logits = tf.stack([model(images, training=False)
                             for _ in range(FLAGS.num_eval_samples)], axis=0)
        else:
          logits = model(images, training=False)
        probs = tf.nn.softmax(logits)
        # print(probs.shape)
        threads.append(threading.Thread(target=compute_entropy_uncertainty, args=(labels, probs, epoch)))
  
        # compute_entropy_uncertainty(labels, probs, epoch)
        # epoch_count += 1
        # exit()
        if FLAGS.num_eval_samples > 1:
          probs = tf.reduce_mean(probs, axis=0)
  
        negative_log_likelihood = tf.reduce_mean(
            tf.keras.losses.sparse_categorical_crossentropy(labels, probs))
  
        if dataset_name == 'clean':
          metrics[f'{dataset_split}/negative_log_likelihood'].update_state(
              negative_log_likelihood)
          metrics[f'{dataset_split}/accuracy'].update_state(labels, probs)
          metrics[f'{dataset_split}/ece'].add_batch(probs, label=labels)
        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)].add_batch(
              probs, label=labels)
  
      for _ in tf.range(tf.cast(num_steps, tf.int32)):
        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):
      threads = []
      logging.info('Starting to run epoch: %s', epoch)
      train_step(train_iterator)
  
      current_step = (epoch + 1) * steps_per_epoch
      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))
      logging.info(message)
  
      if validation_dataset:
        validation_iterator = iter(validation_dataset)
        test_step(
            validation_iterator, 'validation', 'clean', steps_per_validation, epoch, threads)
      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)
        logging.info('Starting to run eval at epoch: %s', epoch)
        test_start_time = time.time()
        test_step(test_iterator, 'test', dataset_name, steps_per_eval, "test", threads)
        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)
      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)
  
      for x in threads: 
          x.start()
  
      for index, thread in enumerate(threads):
          logging.info("Main    : before joining thread %d.", index)
          thread.join()
          logging.info("Main    : thread %d done", index)
  
      # Metrics from Robustness Metrics (like ECE) will return a dict with a
      # single key/value, instead of a scalar.
      total_results = {
          k: (list(v.values())[0] if isinstance(v, dict) else v)
          for k, v in total_results.items()
      }
      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)
  
    # u.batch_uncertainty()
    final_checkpoint_name = checkpoint.save(
        os.path.join(FLAGS.output_dir, 'checkpoint'))
    logging.info('Saved last checkpoint to %s', final_checkpoint_name)
    with summary_writer.as_default():
      hp.hparams({
          'base_learning_rate': FLAGS.base_learning_rate,
          'one_minus_momentum': FLAGS.one_minus_momentum,
          'l2': FLAGS.l2,
          'prior_stddev': FLAGS.prior_stddev,
          'stddev_init': FLAGS.stddev_init,
      })
  
  
  
  
  if __name__ == '__main__':
    app.run(main)
  