<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
Step3:  Hyperparameter Ensembles for Robustness and Uncertainty Quantification
Florian Wenzel, April 8th 2021. Licensed under the Apache License, Version 2.0.
Recently, we proposed Hyper-deep Ensembles (Wenzel et al., NeurIPS 2020) a simple, yet powerful, extension of deep ensembles. The approach works with any given deep network architecture and, therefore, can be easily integrated (and improve) a machine learning system that is already used in production.
Hyper-deep ensembles improve the performance of a given deep network by forming an ensemble over multiple variants of that architecture where each member uses different hyperparameters. In this notebook we consider a ResNet-20 architecture with block-wise $\ell_2$-regularization parameters and a label smoothing parameter. We construct an ensemble of 4 members where each member uses a  different set of hyperparameters. This leads to an ensemble of diverse members, i.e., members that are complementary in their predictions. The final ensemble greatly improves the prediction performance and the robustness of the model, e.g., in out-of-distribution settings.
Let's start with some boilerplate code for data loading and the model definition.
Requirements
Step4:  Let's construct the hyper-deep ensemble over a ResNet-20 architecture
This is the (simplified) hyper-deep ensembles construction pipeline
1. Random search
Step5:  Step 2
Step6:  Now we are ready to construct the ensemble.
* In the first step, we take the best model (on the validation set) -> model_1.
* In the second step, we fix model_1 and try all models in our model pool and construct the ensemble [model_1, model_2]. We select the model model_2 that leads to the highest performance gain.
* In the third step, we fix model_1, model_2 and choose model_3 to construct an ensemble [model_1, model_2, model_3] that leads to the highest performance gain over step 2.
* ... and so on, until the desired ensemble size is reached or no performance gain could be achieved anymore.
Step7:  Evaluation on the test set
Let's see how the hyper-deep ensemble performs on the test set.
Step8:  Here is the final ensemble performance
We gained almost 2 percentage points in terms of accuracy over the best single model!
<ASSISTANT_TASK:>
Python Code:
import tensorflow as tf
import tensorflow_datasets as tfds
import numpy as np
import uncertainty_baselines as ub
def _ensemble_accuracy(labels, logits_list):
  Compute the accuracy resulting from the ensemble prediction.
  per_probs = tf.nn.softmax(logits_list)
  probs = tf.reduce_mean(per_probs, axis=0)
  acc = tf.keras.metrics.SparseCategoricalAccuracy()
  acc.update_state(labels, probs)
  return acc.result()
def _ensemble_cross_entropy(labels, logits):
  logits = tf.convert_to_tensor(logits)
  ensemble_size = float(logits.shape[0])
  labels = tf.cast(labels, tf.int32)
  ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=tf.broadcast_to(labels[tf.newaxis, ...], tf.shape(logits)[:-1]),
      logits=logits)
  nll = -tf.reduce_logsumexp(-ce, axis=0) + tf.math.log(ensemble_size)
  return tf.reduce_mean(nll)
def greedy_selection(val_logits, val_labels, max_ens_size, objective='nll'):
  Greedy procedure from Caruana et al. 2004, with replacement.
  assert_msg = 'Unknown objective type (received {}).'.format(objective)
  assert objective in ('nll', 'acc', 'nll-acc'), assert_msg
  # Objective that should be optimized by the ensemble. Arbitrary objectives,
  # e.g., based on nll, acc or calibration error (or combinations of those) can
  # be used.
  if objective == 'nll':
    get_objective = lambda acc, nll: nll
  elif objective == 'acc':
    get_objective = lambda acc, nll: acc
  else:
    get_objective = lambda acc, nll: nll-acc
  best_acc = 0.
  best_nll = np.inf
  best_objective = np.inf
  ens = []
  def get_ens_size():
    return len(set(ens))
  while get_ens_size() < max_ens_size:
    current_val_logits = [val_logits[model_id] for model_id in ens]
    best_model_id = None
    for model_id, logits in enumerate(val_logits):
      acc = _ensemble_accuracy(val_labels, current_val_logits + [logits])
      nll = _ensemble_cross_entropy(val_labels, current_val_logits + [logits])
      obj = get_objective(acc, nll)
      if obj < best_objective:
        best_acc = acc
        best_nll = nll
        best_objective = obj
        best_model_id = model_id
    if best_model_id is None:
      print('Ensemble could not be improved: Greedy selection stops.')
      break
    ens.append(best_model_id)
  return ens, best_acc, best_nll
def parse_checkpoint_dir(checkpoint_dir):
  Parse directory of checkpoints.
  paths = []
  subdirectories = tf.io.gfile.glob(os.path.join(checkpoint_dir, '*'))
  is_checkpoint = lambda f: ('checkpoint' in f and '.index' in f)
  print('Load checkpoints')
  for subdir in subdirectories:
    for path, _, files in tf.io.gfile.walk(subdir):
      if any(f for f in files if is_checkpoint(f)):
        latest_checkpoint = tf.train.latest_checkpoint(path)
        paths.append(latest_checkpoint)
        print('.', end='')
        break
  print('')
  return paths
DATASET = 'cifar10'
TRAIN_PROPORTION = 0.95
BATCH_SIZE = 64
ENSEMBLE_SIZE = 4
CHECKPOINT_DIR = 'gs://gresearch/reliable-deep-learning/checkpoints/baselines/cifar/hyper_ensemble/'
# Load data.
ds_info = tfds.builder(DATASET).info
num_classes = ds_info.features['label'].num_classes
# Test set.
steps_per_eval = ds_info.splits['test'].num_examples // BATCH_SIZE 
test_dataset = ub.datasets.get(
      DATASET,
      split=tfds.Split.TEST).load(batch_size=BATCH_SIZE)
# Validation set.
validation_percent = 1 - TRAIN_PROPORTION
val_dataset = ub.datasets.get(
    dataset_name=DATASET,
    split=tfds.Split.VALIDATION,
    validation_percent=validation_percent,
    drop_remainder=False).load(batch_size=BATCH_SIZE)
steps_per_val_eval = int(ds_info.splits['train'].num_examples *
                          validation_percent) // BATCH_SIZE
Explanation: Hyperparameter Ensembles for Robustness and Uncertainty Quantification
Florian Wenzel, April 8th 2021. Licensed under the Apache License, Version 2.0.
Recently, we proposed Hyper-deep Ensembles (Wenzel et al., NeurIPS 2020) a simple, yet powerful, extension of deep ensembles. The approach works with any given deep network architecture and, therefore, can be easily integrated (and improve) a machine learning system that is already used in production.
Hyper-deep ensembles improve the performance of a given deep network by forming an ensemble over multiple variants of that architecture where each member uses different hyperparameters. In this notebook we consider a ResNet-20 architecture with block-wise $\ell_2$-regularization parameters and a label smoothing parameter. We construct an ensemble of 4 members where each member uses a  different set of hyperparameters. This leads to an ensemble of diverse members, i.e., members that are complementary in their predictions. The final ensemble greatly improves the prediction performance and the robustness of the model, e.g., in out-of-distribution settings.
Let's start with some boilerplate code for data loading and the model definition.
Requirements:
bash
!pip install "git+https://github.com/google/uncertainty-baselines.git#egg=uncertainty_baselines"
End of explanation
# The model architecture we want to form the ensemble over
# here, we use the original ResNet-20 model by He et al. 2015.
model = ub.models.wide_resnet(
    input_shape=ds_info.features['image'].shape,
    depth=22,
    width_multiplier=1,
    num_classes=num_classes,
    l2=0.,
    version=1)
# Load checkpoints:
# These are 100 checkpoints and loading will take a few minutes.
ensemble_filenames = parse_checkpoint_dir(CHECKPOINT_DIR)
model_pool_size = len(ensemble_filenames)
checkpoint = tf.train.Checkpoint(model=model)
print('Model pool size: {}'.format(model_pool_size))
Explanation: Let's construct the hyper-deep ensemble over a ResNet-20 architecture
This is the (simplified) hyper-deep ensembles construction pipeline
1. Random search: train several models on the train set using different (random) hyperparameters.
2. Ensemble construction: on a validation set using a greedy selection method.
Remark:
In this notebook we use a slightly simplified version of the pipeline compared to the approach of the original paper (where an additional stratification step is used). Additionally, after selecting the optimal hyperparameters the ensemble performance can be improved even more by retraining the selected models on the full train set (i.e., this time not reserving a portion for validation). The simplified pipeline in this notebook is slightly less performant but easier to implement. The simplified pipeline is similar to the ones used by Caranua et al., 2004 and Zaidi et al., 2020 in the context of neural architecture search.
Step 1: Random Hyperparameter Search
We start by training 100 different versions of the ResNet-20 using different $\ell_2$-regularization parameters and label smoothing parameters. Since this would take some time we have already trained the models using a standard training script (which can be found here) and directly load the checkpoints (which can be browsed here).
End of explanation
# Compute the logits on the validation set.
val_logits, val_labels = [], []
for m, ensemble_filename in enumerate(ensemble_filenames):
  # Enforce memory clean-up.
  tf.keras.backend.clear_session()
  checkpoint.restore(ensemble_filename)
  val_iterator = iter(val_dataset)
  val_logits_m = []
  for _ in range(steps_per_val_eval):
    inputs = next(val_iterator)
    features = inputs['features']
    labels = inputs['labels']
    val_logits_m.append(model(features, training=False))
    if m == 0:
      val_labels.append(labels)
  val_logits.append(tf.concat(val_logits_m, axis=0))
  if m == 0:
    val_labels = tf.concat(val_labels, axis=0)
  if m % 10 == 0 or m == model_pool_size - 1:
    percent = (m + 1.) / model_pool_size
    message = ('{:.1%} completion for prediction on validation set: '
                'model {:d}/{:d}.'.format(percent, m + 1, model_pool_size))
    print(message)
Explanation: Step 2: Construction of the hyperparameter ensemble on the validation set
First we compute the logits of all models in our model pool on the validation set.
End of explanation
# Ensemble construction by greedy member selection on the validation set.
selected_members, val_acc, val_nll = greedy_selection(val_logits, val_labels,
                                                        ENSEMBLE_SIZE,
                                                        objective='nll')
unique_selected_members = list(set(selected_members))
message = ('Members selected by greedy procedure: model ids = {} (with {} '
            'unique member(s)).').format(
                selected_members, len(unique_selected_members))
print(message)
Explanation: Now we are ready to construct the ensemble.
* In the first step, we take the best model (on the validation set) -> model_1.
* In the second step, we fix model_1 and try all models in our model pool and construct the ensemble [model_1, model_2]. We select the model model_2 that leads to the highest performance gain.
* In the third step, we fix model_1, model_2 and choose model_3 to construct an ensemble [model_1, model_2, model_3] that leads to the highest performance gain over step 2.
* ... and so on, until the desired ensemble size is reached or no performance gain could be achieved anymore.
End of explanation
# Evaluate the following metrics on the test set.
metrics = {
    'ensemble/negative_log_likelihood': tf.keras.metrics.Mean(),
    'ensemble/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
}
metrics_single = {
    'single/negative_log_likelihood': tf.keras.metrics.SparseCategoricalCrossentropy(),
    'single/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
}
# Compute logits for each ensemble member on the test set.
logits_test = []
for m, member_id in enumerate(unique_selected_members):
  ensemble_filename = ensemble_filenames[member_id]
  checkpoint.restore(ensemble_filename)
  logits = []
  test_iterator = iter(test_dataset)
  for _ in range(steps_per_eval):
    features = next(test_iterator)['features']
    logits.append(model(features, training=False))
  logits_test.append(tf.concat(logits, axis=0))
logits_test = tf.convert_to_tensor(logits_test)
print('Completed computation of member logits on the test set.')
# Compute test metrics.
test_iterator = iter(test_dataset)
for step in range(steps_per_eval):
  labels = next(test_iterator)['labels']
  logits = logits_test[:, (step*BATCH_SIZE):((step+1)*BATCH_SIZE)]
  labels = tf.cast(labels, tf.int32)
  negative_log_likelihood = _ensemble_cross_entropy(labels, logits)
  # Per member output probabilities.
  per_probs = tf.nn.softmax(logits)
  # Ensemble output probabilites.
  probs = tf.reduce_mean(per_probs, axis=0)
  metrics['ensemble/negative_log_likelihood'].update_state(
      negative_log_likelihood)
  metrics['ensemble/accuracy'].update_state(labels, probs)
  # For comparison compute performance of the best single model,
  # this is by definition the first model that was selected by the greedy 
  # selection method.
  logits_single = logits_test[0, (step*BATCH_SIZE):((step+1)*BATCH_SIZE)]
  probs_single = tf.nn.softmax(logits_single)
  metrics_single['single/negative_log_likelihood'].update_state(labels, logits_single)
  metrics_single['single/accuracy'].update_state(labels, probs_single)
  percent = (step + 1) / steps_per_eval
  if step % 25 == 0 or step == steps_per_eval - 1:
    message = ('{:.1%} completion final test prediction'.format(percent))
    print(message)
ensemble_results = {name: metric.result() for name, metric in metrics.items()}
single_results = {name: metric.result() for name, metric in metrics_single.items()}
Explanation: Evaluation on the test set
Let's see how the hyper-deep ensemble performs on the test set.
End of explanation
print('Ensemble performance:')
for m, val in ensemble_results.items():
  print('   {}: {}'.format(m, val))
print('\nFor comparison:')
for m, val in single_results.items():
  print('   {}: {}'.format(m, val))
Explanation: Here is the final ensemble performance
We gained almost 2 percentage points in terms of accuracy over the best single model!
End of explanation
<END_TASK>