#AUTOGENERATED! DO NOT EDIT! File to edit: dev/04_dl.ipynb (unless otherwise specified).

__all__ = ['IMAGENET_STATS', 'gpu_found', 'find_layer_groups', 'inverse_sigmoid', 'mish', 'normalize', 'denormalize',
           'show_batch', 'show_image_variants', 'ImageSequence', 'CosineLRScheduler', 'BatchHistory',
           'PerformanceThreshold', 'plot_history', 'show_predictions', 'hybrid_quantize']

#Cell
import bz2
import cv2
from itertools import chain, zip_longest
import matplotlib.pyplot as plt
import numpy as np
from operator import gt, lt
import os
import pandas as pd
import pickle
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.utils import Sequence
from tensorflow.python.client import device_lib
import warnings

from htools import AutoInit, timeboxed, auto_repr

#Cell
IMAGENET_STATS = np.array([[0.485, 0.456, 0.406],
                           [0.229, 0.224, 0.225]])

#Cell
def gpu_found():
    """Check if any tensorflow can find any gpu's on device.

    Returns
    -------
    bool
    """
    return any('GPU' in str(dev) for dev in device_lib.list_local_devices())

#Cell
def find_layer_groups(model, percent=.5, skip_head=True, plot=True):
    """Tool to help determine model layer groups.

    Parameters
    ----------
    model: tf.keras.Model
    percent: float
        Percent of model weights to find the split point for. For example,
        .3 will find the layer such that 30% of the model weights lie in or
        before that layer.
    skip_head: bool
        Specify whether to exclude the last layer from calculations. We often
        choose to treat the last layer as its own group for transfer learning
        (group 3), so if we want group 1 and 2 to be equally sized in this
        scenario, we must use skip_head=True.
    plot: bool
        If True, plot the number of weights vs. layer number.

    Returns
    -------
    tuple(str, float): Name of layer, cumulative % of model weights contained
        in that layer and its preceding layers.
    """
    name2size = {k: np.prod(v) for k, v in model.layer_dims().items()}
    name2cum = dict(zip(name2size.keys(), np.cumsum(list(name2size.values()))))

    if skip_head:
        total = sorted(name2cum.values())[-2]
    else:
        total = max(name2cum.values())
    cutoff = total * percent

    if plot:
        plt.plot(list(name2cum.values()))
        plt.axhline(cutoff, c='red', ls='--')
        plt.show()

    # Return name of the layer where the cumulative total of weights in the
    # model first surpasses the specified percent. This is approximately where
    # group 2 could start.
    for k, v in name2cum.items():
        if v > cutoff:
            return k, v / total

#Cell
def inverse_sigmoid(y):
    """Use to determine the bias initializer for the final linear layer of
    model.

    Parameters
    ----------
    x: float
        Value between 0 and 1 (e.g. the proportion of the training data that
        are postives).

    Returns
    -------
    float: Inverse sigmoid of input.
        I.E. if y=sigmoid(x), inverse_sigmoid(y)=x.
    """
    return np.log(y / (1-y))

#Cell
def mish(x):
    """Mish: A Self Regularized Non-Monotonic Neural Activation Function

    https://arxiv.org/pdf/1908.08681v1.pdf

    Parameters
    ----------
    x: tf.Tensor[float]
        Input tensor.

    Returns
    -------
    tf.Tensor[float]: Tensor of same shape as input x.
    """
    return x * K.tanh(K.softplus(x))

#Cell
def normalize(img, stats=None):
    """Normalize an image using mean and standard deviation.
    When using pretrained models, the stats should come from the dataset used
    for training, not our data.

    Parameters
    ----------
    img: np.array
        Image or batch of images with channel dimension last.
    stats: np.array
        Channel-wise image stats, where stats[0] contains means and stats[1]
        contains standard deviations. If None (the default), imagenet stats
        will be used.

    Returns
    -------
    """
    # Truth value of array is ambiguous so check explicitly.
    if stats is None:
        stats = IMAGENET_STATS
    return (img - stats[0]) / stats[1]

#Cell
def denormalize(img, stats=None):
    """Reverse normalization process done using mean and standard
    deviation. When using pretrained models, the stats should come from the
    dataset used for training, not our data.

    Parameters
    ----------
    img: np.array
        Image or batch of images with channel dimension last.
    stats: np.array
        Channel-wise image stats, where stats[0] contains means and stats[1]
        contains standard deviations. If None (the default), imagenet stats
        will be used.

    Returns
    -------
    """
    if stats is None:
        stats = IMAGENET_STATS
    return img * stats[1] + stats[0]

#Cell
def show_batch(data_gen, denorm=False, size=(8, 8), max_images=36, **kwargs):
    """
    Parameters
    ----------
    data_gen: DirectoryIterator
        Object returned by ImageDataGenerator.flow_from_directory().
    denorm: bool
        Specify whether denormalization is needed.
    size: tuple[int]
        Figure size.
    max_images: int
        Max number of images to display. When batch sizes are large,
        individual images may be too small to see without either
        increasing size or decreasing max_images.
    kwargs: stats
        If specified, will be passed to denormalize function.
    """
    bs = min(data_gen.batch_size, max_images)
    ncols = int(np.ceil(np.sqrt(bs)))
    i2class = {v: k for k, v in data_gen.class_indices.items()}
    x, y = map(lambda x: x[:bs], next(data_gen))

    # Some generators normalized images w/ imagenet stats.
    if denorm: x = denormalize(x, **kwargs)

    # Create square grid of images. Each image will be labeled with class name.
    fig, ax = plt.subplots(ncols, ncols, figsize=size)
    if ncols == 1: ax = np.array(ax)

    # Must turn off axis for each subplot separately.
    for x, y, ax_ in zip_longest(x, y, ax.flatten()):
        try:
            ax_.imshow(x)
            ax_.set_title(i2class[y])
        except TypeError:
            pass
        ax_.axis('off')

    plt.tight_layout()
    plt.show()

#Cell
def show_image_variants(data_gen, n=4, i=1):
    """Plot n variants of an image from a data generator to get
    a sense of how image augmentation is working. Images are denormalized
    with ImageNet stats before being displayed.

    Parameters
    ----------
    data_gen: iterable
        ImageDataGenerator.flow_from_directory() where each batch is
        the size of the whole dataset. For this reason, it will be
        much faster if constructed from a directory with few images.
    n: int
        Number of variation to display.
    i: int
        Selects i'th image in batch.
    """
    # Plot the images given by the iterator.
    imgs = (denormalize(next(data_gen)[0][i]) for _ in range(n))
    fig, ax = plt.subplots(nrows=i, ncols=n, figsize=(18, 18))
    for ax_i, img in zip(ax, imgs):
        ax_i.imshow(img)
        ax_i.axis('off')
    plt.show()

#Cell
class ImageSequence(Sequence):

    def __init__(self, dir_,  bs, pos_class, image_shape=(224, 224, 3),
                 transforms=None, rescale=1/255, shuffle=True, max_len=None,
                 save_path=None, imagenet_norm=True, seed=1):
        """
        Parameters
        ----------
        dir_: str
            Parent directory. Must contain sub-directories for each class.
        bs: int
            Size of each mini batch.
        pos_class: str
            Name of class to treat as positive (label of 1). Must match the
            name of the directory containing these images.
        image_shape: tuple(int)
            Shape of image with channels last (H x W x C).
        transforms: albumentations.core.composition.Compose
            Compose object from albumentations package. This must be callable
            such that transforms(image=x)['image'] returns the transformed
            image x.
        rescale: float
            Value to multiply image rgb values by. For a standard image
            dataset, the default 1/255 will scale values to lie between 0 and
            1.
        shuffle: bool
            Whether to shuffle the data after each epoch. This will result in
            different mini batches, which may have a regularizing effect.
        max_len: int (optional)
            If specified, will limit the number of samples in the dataset. This
            makes it easier to experiment with subsets if we have a huge number
            of files.
        save_path: str
            Path to save filenames to. Should end with a pickle file.
        imagenet_norm: bool
            If True, normalize dataset with imagenet mean and standard
            deviation.
        seed: int
            Random seed used when shuffling data, if shuffle=True. User may
            pass seed=None to avoid setting a seed.
        """
        if seed is not None:
            np.random.seed(seed)

        self.dir_ = dir_
        self.image_shape = image_shape
        self.bs = bs
        self.pos_class = pos_class
        self.transforms = transforms
        self.rescale = rescale
        self.shuffle = shuffle
        self.save_path = save_path
        self.norm = imagenet_norm
        if self.norm:
            self.means = np.array([0.485, 0.456, 0.406])
            self.stds = np.array([0.229, 0.224, 0.225])

        # Compile list of file paths. Labels will be generated on the fly.
        self.files = self.get_files(max_len)
        self.i2label = self._get_labels()
        self.n_samples = len(self.files)
        self.on_epoch_end()

        if save_path:
            self.save()

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.files)

    def load_img(self, path):
        img = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
        return cv2.resize(img, self.image_shape[:-1])

    def normalize(self, img):
        """Normalize an image using mean and standard deviation. When using
        pretrained models, the stats should come from the dataset used
        for training, not our data.

        Parameters
        ----------
        img: np.array
            Image or batch of images with channel dimension last.
        stats: np.array
            Channel-wise image stats, where stats[0] contains means and
            stats[1] contains standard deviations. If None (the default),
            imagenet stats will be used.

        Returns
        -------
        np.array
        """
        return (img - self.means) / self.stds

    def denormalize(self, img):
        """Reverse normalization process done using mean and standard
        deviation. When using pretrained models, the stats should come from the
        dataset used for training, not our data.

        Parameters
        ----------
        img: np.array
            Image or batch of images with channel dimension last.
        stats: np.array
            Channel-wise image stats, where stats[0] contains means and
            stats[1] contains standard deviations. If None (the default),
            imagenet stats will be used.

        Returns
        -------
        np.array
        """
        return img * self.stds + self.means

    def get_files(self, n):
        """Get list of filenames, discarding any beyond the desired dataset
        size. Shuffling occurs here regardless of self.shuffle choice because
        we want a representative dataset, not all 1 class.
        """
        files = list(self._get_files(self.dir_))
        np.random.shuffle(files)
        return files[:n]

    @staticmethod
    def _get_files(dir_):
        """Get files contained in all subdirectories of a parent directory.

        Parameters
        ----------
        dir_: str
            Name of parent directory. It should contain a subdirectory for each
            class.

        Returns
        -------
        itertools.chain: Generator with all files in subdirectories.
        """
        return chain.from_iterable([
            [f.path for f in os.scandir(arg.path)]
            for arg in os.scandir(dir_)
            if arg.is_dir() and not arg.name.startswith('.')
        ])

    def _get_labels(self):
        """Called automatically on instantiation to create a class
        number to class name mapping in the form of a list.

        Returns
        -------
        List[str]: List of labels, where the positive class is last
            (i.e. in binary classification for classes SFW and NSFW
             where NSFW is positive, we return ['SFW', 'NSFW']).
        """
        labels = [arg.name for arg in os.scandir(self.dir_) if arg.is_dir()]
        return sorted(labels, key=lambda x: x == self.pos_class)

    def show_batch(self, i=None, denorm=True, size=(8, 8)):
        """
        Parameters
        ----------
        i: int
            Batch number to display. If not specified, a random batch is
            chosen.
        denorm: bool
            Specify whether denormalization is needed.
        size: tuple[int]
            Figure size.
        """
        if i is None:
            i = np.random.randint(0, len(self))
        x, y = self[i]
        ncols = int(np.ceil(np.sqrt(self.bs)))

        # Ensure images are displayed correctly and without warning.
        if denorm:
            x = self.denormalize(x)
        x = x.round(5)

        # Plot each image and label with class name.
        fig, ax = plt.subplots(ncols, ncols, figsize=size)
        if ncols == 1:
            ax = np.array(ax)

        # Must turn off axis for each subplot separately.
        for i, (x_, y_, ax_) in enumerate(zip_longest(x, y, ax.flatten())):
            try:
                ax_.imshow(x_)
                ax_.set_title(self.i2label[int(y_)])
            except TypeError:
                pass
            ax_.axis('off')

        plt.tight_layout()
        plt.show()

    def save(self):
        """Pickle and zip generator for later use. Results should be
        reproducible anyway when setting a random seed, but this gives us extra
        insurance.
        """
        print(f'Saving ImageGenerator to {self.save_path}.')
        with bz2.BZ2File(self.save_path, 'w') as f:
            pickle.dump(self, f)

    @staticmethod
    def from_zipfile(path):
        """Load a previously saved ImageSequence.
        """
        with bz2.BZ2File(path, 'r') as f:
            seq = pickle.load(f)
        return seq

    def __len__(self):
        """Returns number of batches per epoch, NOT the number of total images.
        """
        return int(np.ceil(self.n_samples / self.bs))

    def _get_one_xy(self, i):
        """Load a single (image, label) pair.

        Parameters
        ----------
        i: int
            Integer used to index into list of file names.

        Returns
        -------
        tuple[np.array, int]: x, y
        """
        path = self.files[i]
        y = path.split('/')[-2] == self.pos_class
        x = self.load_img(path)
        if self.transforms:
            x = self.transforms(image=x)['image']
        if self.norm:
            x = self.normalize(x * self.rescale)
        return x, y

    def __getitem__(self, i):
        """Retrieve 1 batch.

        Returns
        -------
        np.array, np.array: First item is X, a mini batch of images
            with shape: (batch size, height, width, channels).
            Second item is y, the corresponding labels where 1 indicates NSFW
            and 0 indicates SFW, with shape: (batch size, ).
        """
        # Make sure that we don't select index larger than max idx.
        idx = range(i * self.bs, min(len(self.files), (i+1) * self.bs))
        xy_pairs = [self._get_one_xy(i) for i in idx]
        x, y = map(lambda arg: np.array(arg, dtype=np.float32), zip(*xy_pairs))
        return x, y

    def __eq__(self, obj):
        """Check if two ImageSequence objects are equivalent. Useful for
        debugging and when trying to reproduce results.
        """
        if isinstance(obj, ImageSequence):
            return all(v == getattr(obj, k)
                       for k, v in self._comparable_attrs().items())
        return False

    def __ne__(self, obj):
        """Check if two objects are not equal. Calls __eq__ implicitly.
        """
        return not self == obj

    def _comparable_attrs(self):
        """Helper method for checking equality.
        """
        return {k: v for k, v in self.__dict__.items()
                if k not in ('means', 'stds')}

    def __repr__(self):
        return (f'ImageSequence(dir_={self.dir_}, bs={self.bs}, '
                f'image_shape={self.image_shape}), shuffle={self.shuffle})')

#Cell
class CosineLRScheduler(Callback):
    """Learning rate scheduler that makes updates each batch. Built-in
    Keras version only updates once per epoch.
    """

    def __init__(self, max_lr, epochs, iters_per_epoch, warm=0.3,
                 restarts=False, cycle_len=5, cycle_decay=0.0, min_lr=None,
                 verbose=False):
        """
        Parameters
        ----------
        max_lr: float
            Maximum learning rate to use during training.
        epochs: int
            Number of epochs to train for.
        iters_per_epoch: int
            Number of batches in each epoch.
        warm: float
            Percent of training run (or cycle length) devoted to the increasing
            portion of the schedule. Default 0.3.
        restarts: bool
            Specifies whether to use restarts, i.e. use a cyclical LR.
            True: Version of cosine annealing with restarts. In one
                  cycle, LR starts high and gradually decreases.
                  At the start of the next cycle, it is
                  immediately increased again.
            False: Version of cosine annealing where LR increases
                   for first 30% of training, then decreases for
                   remaining 70%.
        cycle_len: int
            Number of epochs contained in a single cycle. Only used
            when scheduler uses restarts.
        cycle_decay: float
            Scalar to decay the learning rate at the end of each cycle.
            This is only used with restarts, since the regular cosine
            annealing already decays the LR over time.
            E.g. 1.0 will use no decay.
            0.9 means that cycle 2 LRs = cycle 1 LRs * 0.9,
            cycle 3 LRs = cycle 1 LRs * .81,
            etc.
        min_lr: float
            Minimum learning rate. If None is specified, it will be set
            to max_lr / 25.
        """
        super().__init__()
        self.max_lr = max_lr
        self.min_lr = min_lr or max_lr / 25
        self.epochs = epochs
        self.iters_per_epoch = iters_per_epoch
        self.iters = epochs * iters_per_epoch
        self.warm = warm

        if restarts and cycle_len > self.iters_per_epoch:
            warnings.warn('Epoch consist of less than 1 full cycle.')

        self.cycle_len = cycle_len
        self.cycle_decay = cycle_decay
        self.restarts = restarts
        self.lrs = self._schedule(restarts)
        self.curr_iter = 0
        self.verbose = verbose

    def on_batch_begin(self, batch, logs=None):
        """Update learning rate based on current iteration."""
        if self.verbose:
            print(f'\nUpdate LR from {self.model.optimizer.lr.numpy():.3E}',
                  end='')

        # Continue using min_lr if training past specified max iters.
        try:
            lr = self.lrs[self.curr_iter]
        except IndexError:
            lr = self.lrs[-1]
        K.set_value(self.model.optimizer.lr, lr)

        self.curr_iter += 1
        if self.verbose:
            print(f' to {self.model.optimizer.lr.numpy():.3E}')

    @staticmethod
    def _cosine_anneal(iters, lr1, lr2):
        """Helper function for _cosine_tri_lr().

        Parameters
        ----------
        iters: int
            Number of iterations in segment.
        lr1: float
            Learning rate at start of segment.
        lr2: float
            Learning rate at end of segment.

        Returns
        -------
        np.array
        """
        i = np.arange(iters)
        return lr2 + (lr1 - lr2)*(1 + np.cos(np.pi * i/iters))/2

    def _cosine_schedule(self):
        """Cosine annealing scheduler. Computes learning rates for each
        iteration.

        Returns
        -------
        np.array
        """
        seg1 = self._cosine_anneal(int(self.warm * self.iters),
                                   self.min_lr,
                                   self.max_lr)
        seg2 = self._cosine_anneal(int(np.ceil((1 - self.warm) * self.iters)),
                                   self.max_lr,
                                   self.min_lr)
        return np.concatenate((seg1, seg2))

    def _cosine_restarts_schedule(self):
        """Cosine annealing with restarts."""
        cycles = int(np.ceil(self.iters /
                             (self.cycle_len * self.iters_per_epoch)))
        cycle_iters = self.cycle_len * self.iters_per_epoch
        lrs = [self._cosine_anneal(cycle_iters, self.max_lr, self.min_lr)
               / (1 + self.cycle_decay * i) for i in range(cycles)]
        return np.concatenate(lrs)

    def _schedule(self, restarts):
        """Wrapper to schedule learning rates depending on chosen method.

        Parameters
        ----------
        restarts: bool
            If True, use schedule with restarts. If False, use regular
            cosine annealing that spans whole duration of training.

        Returns
        -------
        np.array: LR for each iteration (i.e. output[i] is the LR to use
            at iteration i).
        """
        if restarts:
            return self._cosine_restarts_schedule()
        return self._cosine_schedule()

    def plot_lrs(self, path=None):
        """Display learning rate by iteration.

        Note: If the plot is not as smooth as expected, this likely
        means that there are very few iterations per epoch
        (i.e. the batch size is very large, at least in relative terms).
        """
        plt.plot(self.lrs)
        plt.xlabel('Iteration')
        plt.ylabel('Learning Rate')
        plt.title('Learning Rate Schedule')
        if path:
            plt.savefig(path)
            plt.close()
        else:
            plt.show()

    def reset(self):
        """Reset the current iteration to zero. Note that if the learning rate
        has been decayed, it will now be back at the initial value.
        """
        self.curr_iter = 0

    def __repr__(self):
        return (f'CosineLRScheduler(max_lr={self.max_lr}, '
                f'epochs={self.epochs}, '
                f'iters_per_epoch={self.iters_per_epoch}, '
                f'restarts={self.restarts}, '
                f'cycle_len={self.cycle_len}, '
                f'min_lr={self.min_lr}, '
                f'verbose={self.verbose})')

#Cell
class BatchHistory(Callback):
    """Callback to store training loss at each mini batch. Built in
    history callback only stores epoch losses. Validation set is only
    evaluated at the end of each epoch so it's not included here.
    """

    def __init__(self, path):
        """Parameters
        -------------
        path: str
            Path to save csv of batch losses to.
        """
        super().__init__()
        self.path = path
        self.losses = []

    def on_batch_end(self, batch, logs=None):
        self.losses.append(logs.get('loss'))

    def on_train_end(self, logs=None):
        with open(self.path, 'a') as f:
            f.write('\n'.join(str(loss) for loss in self.losses) + '\n')

    def plot_batch_losses(self, path=None, smooth=25):
        df = pd.DataFrame(self.losses, columns=['loss'])
        df['smoothed'] = df.loss.ewm(span=smooth).mean()

        # Plot losses and smoothed losses.
        plt.plot(df.loss, lw=.5, c='blue', alpha=.25)
        plt.plot(df.smoothed, lw=.5, c='blue', alpha=1)
        plt.xlabel('Iteration')
        plt.ylabel('Training Loss')
        plt.title('Training Loss by Mini Batch.')
        if path:
            plt.savefig(path)
            plt.close()
        else:
            plt.show()

#Cell
class PerformanceThreshold(Callback):
    """Callback that will stop training if performance falls short of
    some user-specified threshold. This can help avoid wasting time on
    runs that have clearly gone wrong - e.g. learning rate is too large
    and validation BCE loss has ballooned to >2.0 on a balanced dataset.
    """

    def __init__(self, metric='val_loss', threshold=1.0, goal='min'):
        """
        Parameters
        ----------
        metric: str
            Name of metric to monitor.
        threshold: float
            Performance threshold. If the quantity being monitored is
            something we want to minimize, this is the highest
            acceptable value to continue training. If it's a metric to
            maximize like accuracy, this is the minimum acceptable
            value.
        goal: str
            Goal of the quantity being monitored (i.e. are we trying to
            minimize metric or maximize it?). One of ('max', 'min').
        """
        super().__init__()
        assert goal in ('max', 'min')

        self.metric = metric
        self.threshold = threshold
        if goal == 'max':
            self.op = gt
        elif goal == 'min':
            self.op = lt

    def on_epoch_end(self, batch, logs=None):
        curr_val = logs.get(self.metric)
        if curr_val is not None and self.op(self.threshold, curr_val):
            print(f'{self.metric} of {curr_val:.3e} did not satisfy '
                  f'performance threshold of {self.threshold}. '
                   'Stopping training.')
            self.model.stop_training = True

#Cell
def plot_history(*args, vlines=False, metrics=None, path=None):
    """Plot metrics from one or more training run.

    Parameters
    ----------
    args: tensorflow.keras.callbacks.history
        One or more history objects returned from training. If multiple history
        objects are passed in, their stats will be concatenated.
    vlines: bool
        If True, add vertical lines to the plots showing where each training
        run ended.
    metrics: list[str]
        Names of metrics to plot. If excluded, all metrics will be plotted
    """
    if not metrics:
        metrics = {metric.replace('val_', '')
                   for metric in args[0].history.keys()}

    # Concatenate metric lists if multiple history objects passed in.
    hist = {m: list(chain.from_iterable([arg.history[m] for arg in args]))
            for m in args[0].history.keys()}
    train_lens = np.cumsum([len(arg.history['loss']) for arg in args])

    # Create subplot for each metric.
    length = len(metrics)
    fig, ax = plt.subplots(length, 1, figsize=(8, length*2.5))
    for i, m in enumerate(metrics):
        ax[i].set_title(m.title())
        ax[i].plot(hist[m], label=m)

        # Plot validation metric if present.
        val_metric = 'val_' + m
        try:
            ax[i].plot(hist[val_metric], label=val_metric)
        except KeyError:
            pass

        ax[i].legend()
        if not vlines:
            continue

        # Show where each training run ended.
        for length in train_lens[:-1]:
            ax[i].axvline(length, lw=0.8, alpha=0.9, ls='--', c='grey')

    plt.subplots_adjust(hspace=.4)
    if path:
        plt.savefig(path)
        plt.close()
    else:
        plt.show()

#Cell
def show_predictions(df, dg, mode, n=9, sort=True, most_wrong=True,
                     figsize=(8, 8), threshold=None, y_true='y',
                     y_pred='y_pred', y_proba='y_proba'):
    """Display images based on the model's predictions. You can view false
    positives, false negatives, or a random subset of all images.

    Parameters
    ----------
    df: pd.DataFrame
        Labels and predictions df produced by evaluate_model.py script.
    dg: ImageSequence
        Same data generator that was used to produce the predictions in df.
        Because shuffle=False in evaluation script, row i in df corresponds to
        file i in dg.
    mode: str
        'fp' to show false positives, 'fn' to show false negatives, 'any'
        to show any predictions (both correct and incorrect).
    n: int
        Number of images to display.
    sort: bool
        Whether to sort predictions by distance from true label before
        displaying. I.e. if True, will show top 5 "most wrong" predictions.
        If False, shows first 5 wrong predictions.
    most_wrong: bool
        If True, show the predictions where the model was most wrong. If False,
        show the "least wrong" errors (model was wrong but nearly predicted
        the correct class). To get a random selection of images, set
        sort=False.
    threshold: float or None
        If specified, will re-compute hard predictions using this as the
        decision threshold. If None, df must already have column of hard
        predictions.
    hard_col: str
        Name of column in df containing hard predictions.
    soft_col: str
        Name of column in df containing soft predictions.
    """
    if threshold: df = update_hard_preds(df, threshold, soft_col)

    if mode == 'fp':
        df = df[df[y_true] < df[y_pred]]
    elif mode == 'fn':
        df = df[df[y_true] > df[y_pred]]
    elif mode != 'any':
        raise ValueError('Parameter `mode` must be "fp", "fn" or "any".')

    # Sort after filtering hard labels since we may use very high thresholds
    # at times, which could lead to correctly predicted labels that have large
    #  distances between predicted prob and true label.
    ascending = not most_wrong
    if sort:
        df = df.lambda_sort(lambda x: abs(x[y_true] - x[y_proba]),
                            ascending=ascending)
    else:
        df = df.sample(frac=1, replace=False)

    ncols = int(np.ceil(np.sqrt(n)))
    fig, ax = plt.subplots(ncols, ncols, figsize=figsize)
    for idx, ax_ in zip_longest(df.index[:n], ax.flatten()):
        try:
            path = dg.files[idx]
            img = dg.load_img(path)
            ax_.imshow(img)
            title = (f'{path.split("/")[-1]}\n '
                     f'{df.loc[idx, y_proba]:.4f}')
            ax_.set_title(title[-30:], size=10)
        except:
            pass
        ax_.axis('off')

    plt.tight_layout()
    plt.show()

#Cell
def hybrid_quantize(model, path=None):
    """Perform post-training hybrid quantization. Weights will be
    changed from float32 to int8, while biases and activation remain as
    float32. TF claims that a 75% reduction in model size is typical, with
    small effects on accuracy.

    Parameters
    ----------
    model: tf.keras.models.Model
    path: str
        Optionally specify an output path to save the quantized weights to.

    Returns
    -------
    bytes or None: If no path is provided, returns quantized weights. If
        a path is provided, save model and return None.
    """
    # Reduced float quantization.
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
    quantized_bytes = converter.convert()

    # Save quantized weights if requested.
    if path:
        with open(path, 'wb') as f:
            f.write(quantized_bytes)
            print(f'Quantized bytes saved to {path}.')
            return
    return quantized_bytes