<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Run a pre-trained Epinet on ImageNet
This demo shows how to run and evaluate a pre-trained Epinet on ImageNet. Epinet is a new ENN architecture that can supplement any conventional NN and be trained to estimate uncertainty.
An epinet is a neural network with privileged access to inputs and outputs of activation units in the base network.
A subset of these inputs and outputs, denoted by $\phi_\zeta(x)$, are taken as input to the epinet along with an epistemic index $z$.
For epinet parameters $\eta$, the epinet outputs $\sigma_\eta(\phi_\zeta(x), z)$.
To produce an ENN, the output of the epinet is added to that of the base network, though with a "stop gradient" written $[[\cdot]]$
Step1:  Imports
Step2:  Load ImageNet dataset
Our enn library provides functionalities in enn/datasets to load ImageNet, CIFAR10/100, and MNIST datasets. To load these datasets, you need to download that dataset into the default tensorflow dataset directory of ~/tensorflow_datasets/downloads/manual/. 
In this colab, we want to evaluate Epinet on only one small batch of ImageNet test images. To this end, we provide a sample batch of size 100 at https
Step3:  Define a set of evaluation metrics
Our enn library provides the set of known metrics for evaluating the performance of neural networks. These metrics which can be access from enn/metrics can be divided in three categories
Step4:  Load pre-trained Epinet
Pre-trained Epinet can be accessed from ImagenetModels in enn.checkpointing.catalog.py. As of now, we provide pre-trained Epinet based on ResNet-50, ResNet-101, ResNet-152, and ResNet-200. In this colab, we want to load Epinet based on ResNet-50 which can be accessed from the checkpoint RESNET_50_FINAL_EPINET.
Step5:  From the checkpoint, we can recover an enn sampler, which is a function that takes a batch of images and one random key, and returns multiple sample logits per input image. To recover the enn sampler, we can use make_epinet_sampler_from_checkpoint (from enn/checkpoints/utils.py) which takes the checkpoint and also the number of sample logits we want per image (num_enn_samples).
Step6:  Load pre-trained ResNet
To have a better sense of how amazing Epinet is, we can compare its performance with a pretrained ResNet-50.
Pre-trained ResNets can be accessed from ImagenetModels in enn.checkpointing.catalog.py. As of now, we provide pre-trained ResNet-50, ResNet-101, ResNet-152, and ResNet-200.
Step7:  From the checkpoint, we can recover an enn sampler, which is a function that takes a batch of images and one random key, and returns multiple sample logits per input image. To recover the enn sampler for ResNet-50, we can use make_enn_sampler_from_checkpoint (from enn/checkpoints/utils.py) which takes the checkpoint and also the number of sample logits we want per image (num_enn_samples). Here we set num_enn_samples=1, as having num_enn_samples &gt; 1 just results in multiple similar sample logits per input image.
Step8:  Compare Epinet and ResNet results
<ASSISTANT_TASK:>
Python Code:
# Copyright 2022 DeepMind Technologies Limited. 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.
# ==============================================================================
!git clone https://github.com/deepmind/enn.git
!pip install -q enn/
Explanation: Run a pre-trained Epinet on ImageNet
This demo shows how to run and evaluate a pre-trained Epinet on ImageNet. Epinet is a new ENN architecture that can supplement any conventional NN and be trained to estimate uncertainty.
An epinet is a neural network with privileged access to inputs and outputs of activation units in the base network.
A subset of these inputs and outputs, denoted by $\phi_\zeta(x)$, are taken as input to the epinet along with an epistemic index $z$.
For epinet parameters $\eta$, the epinet outputs $\sigma_\eta(\phi_\zeta(x), z)$.
To produce an ENN, the output of the epinet is added to that of the base network, though with a "stop gradient" written $[[\cdot]]$:
$$ f_\theta(x, z) = \mu_\zeta(x) + \sigma_\eta([[\phi_\zeta(x)]], z). $$
We can visualize this network architecture:
For more details about Epinet, refer to the paper 
Epistemic Neural Networks (Osband et al., 2022).
It's recommended to use Runtime-&gt;Change Runtime Type to pick a GPU for speed.
End of explanation
#@title General imports
import warnings
warnings.filterwarnings('ignore')
#@title Development imports
from typing import Callable, NamedTuple
import numpy as np
import pandas as pd
import plotnine as gg
from acme.utils.loggers.terminal import TerminalLogger
import dataclasses
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import optax
import dill
#@title ENN imports
import enn
from enn import datasets
from enn.checkpoints import base as checkpoint_base
from enn.networks.epinet import base as epinet_base
from enn.checkpoints import utils
from enn.checkpoints import imagenet
from enn.checkpoints import catalog
from enn import metrics as enn_metrics
Explanation: Imports
End of explanation
!wget https://storage.googleapis.com/dm-enn/processed_batch.npzs --no-check-certificate
with open('processed_batch.npzs', 'rb') as file:
  batch = dill.load(file)
images, labels = batch['images'], batch['labels']
Explanation: Load ImageNet dataset
Our enn library provides functionalities in enn/datasets to load ImageNet, CIFAR10/100, and MNIST datasets. To load these datasets, you need to download that dataset into the default tensorflow dataset directory of ~/tensorflow_datasets/downloads/manual/. 
In this colab, we want to evaluate Epinet on only one small batch of ImageNet test images. To this end, we provide a sample batch of size 100 at https://storage.googleapis.com/dm-enn/processed_batch.npzs which can be download as follows.
End of explanation
# Define a dict of metrics including `accuracy`, `marginal nll`, and `joint nll`.
evaluation_metrics = {
      'accuracy': enn_metrics.make_accuracy_calculator(),
      'marginal nll': enn_metrics.make_nll_marginal_calculator(),
      'joint nll': enn_metrics.make_nll_polyadic_calculator(tau=10, kappa=2),
  }
Explanation: Define a set of evaluation metrics
Our enn library provides the set of known metrics for evaluating the performance of neural networks. These metrics which can be access from enn/metrics can be divided in three categories:
Marginal: includes metrics like accuracy and marginal negative log-likelihood (NLL) for evaluating marginal predictions.
Joint: includes metrics for evaluating joint predictions. 
Calibration: includes metrics for calculating calibration error.
Each metric takes logits and lables with the following shapes:
  - logits: [num_enn_samples, batch_size, num_classes]
  - labels: [batch_size, 1]
num_enn_samples specifies the number of sample logits per input image.
End of explanation
# Get the Epinet checkpoint
epinet_resnet50_imagenet_ckpt = catalog.ImagenetModels.RESNET_50_FINAL_EPINET.value
epinet_resnet50_imagenet_ckpt
Explanation: Load pre-trained Epinet
Pre-trained Epinet can be accessed from ImagenetModels in enn.checkpointing.catalog.py. As of now, we provide pre-trained Epinet based on ResNet-50, ResNet-101, ResNet-152, and ResNet-200. In this colab, we want to load Epinet based on ResNet-50 which can be accessed from the checkpoint RESNET_50_FINAL_EPINET.
End of explanation
# Set the number of sample logits per input image
num_enn_samples = 100
# Recover the enn sampler
epinet_enn_sampler = utils.make_epinet_sampler_from_checkpoint(
    epinet_resnet50_imagenet_ckpt,
    num_enn_samples=num_enn_samples,)
# Get the epinet logits
key = jax.random.PRNGKey(seed=0)
epinet_logits = epinet_enn_sampler(images, key)
# epinet logits has shape [num_enn_sample, eval_batch_size, num_classes]
epinet_logits.shape
# Labels loaded from our dataset has shape [eval_batch_size,]. Our evaluation
# metrics requires labels to have shape [eval_batch_size, 1].
eval_labels = labels[:, None]
# Evaluate 
epinet_results = {key: float(metric(epinet_logits, eval_labels)) 
                      for key, metric in evaluation_metrics.items()}
epinet_results
Explanation: From the checkpoint, we can recover an enn sampler, which is a function that takes a batch of images and one random key, and returns multiple sample logits per input image. To recover the enn sampler, we can use make_epinet_sampler_from_checkpoint (from enn/checkpoints/utils.py) which takes the checkpoint and also the number of sample logits we want per image (num_enn_samples).
End of explanation
# Get the ResNet-50 checkpoint
resnet50_imagenet_ckpt = catalog.ImagenetModels.RESNET_50.value
resnet50_imagenet_ckpt
Explanation: Load pre-trained ResNet
To have a better sense of how amazing Epinet is, we can compare its performance with a pretrained ResNet-50.
Pre-trained ResNets can be accessed from ImagenetModels in enn.checkpointing.catalog.py. As of now, we provide pre-trained ResNet-50, ResNet-101, ResNet-152, and ResNet-200.
End of explanation
# Set the number of sample logits per input image to 1
num_enn_samples = 1
# Recover the enn sampler
resnet50_enn_sampler = utils.make_enn_sampler_from_checkpoint(
    resnet50_imagenet_ckpt,
    num_enn_samples=num_enn_samples,)
# Get the epinet logits
key = jax.random.PRNGKey(seed=0)
resnet50_logits = resnet50_enn_sampler(images, key)
# ResNet logits has shape [num_enn_sample, eval_batch_size, num_classes]
resnet50_logits.shape
# Labels loaded from our dataset has shape [eval_batch_size,]. Our evaluation
# metrics requires labels to have shape [eval_batch_size, 1].
eval_labels = labels[:, None]
# Evaluate 
resnet50_results = {key: float(metric(resnet50_logits, eval_labels)) 
                      for key, metric in evaluation_metrics.items()}
resnet50_results
Explanation: From the checkpoint, we can recover an enn sampler, which is a function that takes a batch of images and one random key, and returns multiple sample logits per input image. To recover the enn sampler for ResNet-50, we can use make_enn_sampler_from_checkpoint (from enn/checkpoints/utils.py) which takes the checkpoint and also the number of sample logits we want per image (num_enn_samples). Here we set num_enn_samples=1, as having num_enn_samples &gt; 1 just results in multiple similar sample logits per input image.
End of explanation
# Make a dataframe of the results
resnet50_results['model'] = 'resnet'
epinet_results['model'] = 'epinet'
df = pd.DataFrame([resnet50_results, epinet_results])
df
# Compare the results
plt_df = pd.melt(df, id_vars=['model'], value_vars=evaluation_metrics.keys())
p = (gg.ggplot(plt_df)
    + gg.aes(x='model', y='value', fill='model') 
    + gg.geom_col()
    + gg.facet_wrap('variable', scales='free',)
    + gg.theme(figure_size=(14, 4), panel_spacing=0.7)
    )
p
Explanation: Compare Epinet and ResNet results
End of explanation
<END_TASK>