# Programming Assignment

## RealNVP for the LSUN bedroom dataset

### Instructions

In this notebook, you will develop the RealNVP normalising flow architecture from scratch, including the affine coupling layers, checkerboard and channel-wise masking, and combining into a multiscale architecture. You will train the normalising flow on a subset of the LSUN bedroom dataset.

Some code cells are provided for you in the notebook. You should avoid editing provided code, and make sure to execute the cells in order to avoid unexpected errors. Some cells begin with the line:

`#### GRADED CELL ####`

Don't move or edit this first line - this is what the automatic grader looks for to recognise graded cells. These cells require you to write your own code to complete them, and are automatically graded when you submit the notebook. Don't edit the function name or signature provided in these cells, otherwise the automatic grader might not function properly.

### How to submit

Complete all the tasks you are asked for in the worksheet. When you have finished and are happy with your code, press the **Submit Assignment** button at the top of this notebook.

### Let's get started!

We'll start running some imports, and loading the dataset. Do not edit the existing imports in the following cell. If you would like to make further Tensorflow imports, you should add them here.

```python
#### PACKAGE IMPORTS ####

# Run this cell first to import all required packages. Do not make any imports elsewhere in the notebook

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import matplotlib.pyplot as plt

from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Conv2D, BatchNormalization
from tensorflow.keras.optimizers import Adam

tfd = tfp.distributions
tfb = tfp.bijectors

# If you would like to make further imports from tensorflow, add them here

from tensorflow.keras import layers
from tensorflow.keras import regularizers
from tensorflow.keras.layers import Add
```

#### The LSUN Bedroom Dataset

In this assignment, you will use a subset of the [LSUN dataset](https://www.yf.io/p/lsun). This is a large-scale image dataset with 10 scene and 20 object categories. A subset of the LSUN bedroom dataset has been provided, and has already been downsampled and preprocessed into smaller, fixed-size images.

* F. Yu, A. Seff, Y. Zhang, S. Song, T. Funkhouser and J. Xia. "LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop". [arXiv:1506.03365](https://arxiv.org/abs/1506.03365), 10 Jun 2015

Your goal is to develop the RealNVP normalising flow architecture using bijector subclassing, and use it to train a generative model of the LSUN bedroom data subset. For full details on the RealNVP model, refer to the original paper:

* L. Dinh, J. Sohl-Dickstein and S. Bengio. "Density estimation using Real NVP". [arXiv:1605.08803](https://arxiv.org/abs/1605.08803), 27 Feb 2017.

#### Load the dataset

The following functions will be useful for loading and preprocessing the dataset. The subset you will use for this assignment consists of 10,000 training images, 1000 validation images and 1000 test images.

The images have been downsampled to 32 x 32 x 3 in order to simplify the training process.

```python
# Functions for loading and preprocessing the images

def load_image(img):
    img = tf.image.random_flip_left_right(img)
    return img, img

def load_dataset(split):
    train_list_ds = tf.data.Dataset.from_tensor_slices(np.load('./data/{}.npy'.format(split)))
    train_ds = train_list_ds.map(load_image)
    return train_ds
```

```python
# Load the training, validation and testing datasets splits

train_ds = load_dataset('train')
val_ds = load_dataset('val')
test_ds = load_dataset('test')
```

```python
# Shuffle the datasets

shuffle_buffer_size = 1000
train_ds = train_ds.shuffle(shuffle_buffer_size)
val_ds = val_ds.shuffle(shuffle_buffer_size)
test_ds = test_ds.shuffle(shuffle_buffer_size)
```

```python
# Display a few examples

n_img = 4
f, axs = plt.subplots(n_img, n_img, figsize=(14, 14))

for k, image in enumerate(train_ds.take(n_img**2)):
    i = k // n_img
    j = k % n_img
    axs[i, j].imshow(image[0])
    axs[i, j].axis('off')
f.subplots_adjust(wspace=0.01, hspace=0.03)
```

```python
# Batch the Dataset objects

batch_size = 64
train_ds = train_ds.batch(batch_size)
val_ds = val_ds.batch(batch_size)
test_ds = test_ds.batch(batch_size)
```

### Affine coupling layer

We will begin the development of the RealNVP architecture with the core bijector that is called the _affine coupling layer_. This bijector can be described as follows: suppose that $x$ is a $D$-dimensional input, and let $d<D$. Then the output $y$ of the affine coupling layer is given by the following equations:

$$
\begin{align}
y_{1:d} &= x_{1:d} \tag{1}\\
    y_{d+1:D} &= x_{d+1:D}\odot \exp(s(x_{1:d})) + t(x_{1:d}), \tag{2}
\end{align}
$$

where $s$ and $t$ are functions from $\mathbb{R}^d\rightarrow\mathbb{R}^{D-d}$, and define the log-scale and shift operations on the vector $x_{d+1:D}$ respectively.

The log of the Jacobian determinant for this layer is given by $\sum_{j}s(x_{1:d})_j$.

The inverse operation can be easily computed as

$$
\begin{align}
x_{1:d} &= y_{1:d}\tag{3}\\
x_{d+1:D} &= \left(y_{d+1:D} - t(y_{1:d})\right)\odot \exp(-s(y_{1:d})),\tag{4}
\end{align}
$$

In practice, we will implement equations $(1)$ and $(2)$ using a binary mask $b$:

$$
\begin{align}
\text{Forward pass:}\qquad y &= b\odot x + (1-b)\odot\left(x\odot\exp(s(b\odot x)) + t(b\odot x)\right),\tag{5}\\
\text{Inverse pass:}\qquad x &= b\odot y + (1-b)\odot\left(y - t(b\odot x)) \odot\exp( -s(b\odot x)\right).\tag{6}
\end{align}
$$

Our inputs $x$ will be a batch of 3-dimensional Tensors with `height`, `width` and `channels` dimensions. As in the original architecture, we will use both spatial 'checkerboard' masks and channel-wise masks:

Figure 1. Spatial checkerboard mask (left) and channel-wise mask (right). From the original paper.

#### Custom model for log-scale and shift

You should now create a custom model for the shift and log-scale parameters that are used in the affine coupling layer bijector. We will use a convolutional residual network, with two residual blocks and a final convolutional layer. Using the functional API, build the model according to the following specifications:

* The function takes the `input_shape` and `filters` as arguments
* The model should use the `input_shape` in the function argument to set the shape in the Input layer (call this layer `h0`).
* The first hidden layer should be a Conv2D layer with number of filters set by the `filters` argument, and a ReLU activation
* The second hidden layer should be a BatchNormalization layer
* The third hidden layer should be a Conv2D layer with the same number of filters as the input `h0` to the model, and a ReLU activation
* The fourth hidden layer should be a BatchNormalization layer
* The fifth hidden layer should be the sum of the fourth hidden layer output and the inputs `h0`. Call this layer `h1`
* The sixth hidden layer should be a Conv2D layer with filters set by the `filters` argument, and a ReLU activation
* The seventh hidden layer should be a BatchNormalization layer
* The eighth hidden layer should be a Conv2D layer with the same number of filters as `h1` (and `h0`), and a ReLU activation
* The ninth hidden layer should be a BatchNormalization layer
* The tenth hidden layer should be the sum of the ninth hidden layer output and `h1`
* The eleventh hidden layer should be a Conv2D layer with the number of filters equal to twice the number of channels of the model input, and a linear activation. Call this layer `h2`
* The twelfth hidden layer should split `h2` into two equal-sized Tensors along the final channel axis. These two Tensors are the shift and log-scale Tensors, and should each have the same shape as the model input
* The final layer should then apply the `tanh` nonlinearity to the log_scale Tensor. The outputs to the model should then be the list of Tensors `[shift, log_scale]`

All Conv2D layers should use a 3x3 kernel size, `"SAME"` padding and an $l2$ kernel regularizer with regularisation coefficient of `5e-5`.

_Hint: use_ `tf.split` _with arguments_ `num_or_size_splits=2, axis=-1` _to create the output Tensors_.

In total, the network should have 14 layers (including the `Input` layer).

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_conv_resnet(input_shape, filters):
    """
    This function should build a CNN ResNet model according to the above specification,
    using the functional API. The function takes input_shape as an argument, which should be
    used to specify the shape in the Input layer, as well as a filters argument, which
    should be used to specify the number of filters in (some of) the convolutional layers.
    Your function should return the model.

    The function takes the input_shape and filters as arguments
    The model should use the input_shape in the function argument to set the shape in the Input layer (call this layer h0).
    The first hidden layer should be a Conv2D layer with number of filters set by the filters argument, and a ReLU activation
    The second hidden layer should be a BatchNormalization layer
    The third hidden layer should be a Conv2D layer with the same number of filters as the input h0 to the model,
    and a ReLU activation
    The fourth hidden layer should be a BatchNormalization layer
    The fifth hidden layer should be the sum of the fourth hidden layer output and the inputs h0. Call this layer h1
    The sixth hidden layer should be a Conv2D layer with filters set by the filters argument, and a ReLU activation
    The seventh hidden layer should be a BatchNormalization layer
    The eighth hidden layer should be a Conv2D layer with the same number of filters as h1 (and h0), and a ReLU activation
    The ninth hidden layer should be a BatchNormalization layer
    The tenth hidden layer should be the sum of the ninth hidden layer output and h1
    The eleventh hidden layer should be a Conv2D layer with the number of filters equal to
    twice the number of channels of the model input, and a linear activation. Call this layer h2
    The twelfth hidden layer should split h2 into two equal-sized Tensors along the final channel axis.
    These two Tensors are the shift and log-scale Tensors, and should each have the same shape as the model input
    The final layer should then apply the tanh nonlinearity to the log_scale Tensor.
    The outputs to the model should then be the list of Tensors [shift, log_scale]
    """
    h0 = Input(shape = input_shape)
    h = Conv2D(filters = filters,activation = 'relu',kernel_size = (3,3),padding='SAME',kernel_regularizer=regularizers.l2(5e-5))(h0)
    h = BatchNormalization()(h)
    h = Conv2D(filters = input_shape[-1],activation = 'relu',kernel_size = (3,3),padding='SAME',kernel_regularizer=regularizers.l2(5e-5))(h)
    h = BatchNormalization()(h)
    # h1 = Add()([h, h0])
    h1 = tf.add(h, h0)
    h = Conv2D(filters = filters,activation = 'relu',kernel_size = (3,3),padding='SAME',kernel_regularizer=regularizers.l2(5e-5))(h1)
    h = BatchNormalization()(h)
    h = Conv2D(filters = input_shape[-1],activation = 'relu',kernel_size = (3,3),padding='SAME',kernel_regularizer=regularizers.l2(5e-5))(h)
    h = BatchNormalization()(h)
    # h = Add()([h1, h])
    h = tf.add(h1, h)
    h2 = Conv2D(filters = 2*input_shape[-1],activation = 'linear',kernel_size = (3,3),padding='SAME',kernel_regularizer=regularizers.l2(5e-5))(h)
    shift, log_scale = tf.split(h2, num_or_size_splits=2, axis = -1)
    log_scale = tf.math.tanh(log_scale)
    outputs = [shift, log_scale]

    model = tf.keras.Model(inputs=h0, outputs=outputs)

    return model

```

```python
# Test your function and print the model summary

conv_resnet = get_conv_resnet((32, 32, 3), 32)
conv_resnet.summary()
```

    Model: "model_1"
    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_2 (InputLayer)            [(None, 32, 32, 3)]  0
    __________________________________________________________________________________________________
    conv2d_5 (Conv2D)               (None, 32, 32, 32)   896         input_2[0][0]
    __________________________________________________________________________________________________
    batch_normalization_4 (BatchNor (None, 32, 32, 32)   128         conv2d_5[0][0]
    __________________________________________________________________________________________________
    conv2d_6 (Conv2D)               (None, 32, 32, 3)    867         batch_normalization_4[0][0]
    __________________________________________________________________________________________________
    batch_normalization_5 (BatchNor (None, 32, 32, 3)    12          conv2d_6[0][0]
    __________________________________________________________________________________________________
    tf_op_layer_Add_2 (TensorFlowOp [(None, 32, 32, 3)]  0           batch_normalization_5[0][0]
                                                                     input_2[0][0]
    __________________________________________________________________________________________________
    conv2d_7 (Conv2D)               (None, 32, 32, 32)   896         tf_op_layer_Add_2[0][0]
    __________________________________________________________________________________________________
    batch_normalization_6 (BatchNor (None, 32, 32, 32)   128         conv2d_7[0][0]
    __________________________________________________________________________________________________
    conv2d_8 (Conv2D)               (None, 32, 32, 3)    867         batch_normalization_6[0][0]
    __________________________________________________________________________________________________
    batch_normalization_7 (BatchNor (None, 32, 32, 3)    12          conv2d_8[0][0]
    __________________________________________________________________________________________________
    tf_op_layer_Add_3 (TensorFlowOp [(None, 32, 32, 3)]  0           tf_op_layer_Add_2[0][0]
                                                                     batch_normalization_7[0][0]
    __________________________________________________________________________________________________
    conv2d_9 (Conv2D)               (None, 32, 32, 6)    168         tf_op_layer_Add_3[0][0]
    __________________________________________________________________________________________________
    tf_op_layer_split_1 (TensorFlow [(None, 32, 32, 3),  0           conv2d_9[0][0]
    __________________________________________________________________________________________________
    tf_op_layer_Tanh_1 (TensorFlowO [(None, 32, 32, 3)]  0           tf_op_layer_split_1[0][1]
    ==================================================================================================
    Total params: 3,974
    Trainable params: 3,834
    Non-trainable params: 140
    __________________________________________________________________________________________________

You can also inspect your model architecture graphically by running the following cell. It should look something like the following:

```python
# Plot the model graph

tf.keras.utils.plot_model(conv_resnet, show_layer_names=False, rankdir='LR')
```

```python
# Check the output shapes are as expected

print(conv_resnet(tf.random.normal((1, 32, 32, 3)))[0].shape)
print(conv_resnet(tf.random.normal((1, 32, 32, 3)))[1].shape)
```

    (1, 32, 32, 3)
    (1, 32, 32, 3)

#### Binary masks

Now that you have a shift and log-scale model built, we will now implement the affine coupling layer. We will first need functions to create the binary masks $b$ as described above. The following function creates the spatial 'checkerboard' mask.

It takes a rank-2 `shape` as input, which correspond to the `height` and `width` dimensions, as well as an `orientation` argument (an integer equal to `0` or `1`) that determines which way round the zeros and ones are entered into the Tensor.

```python
# Function to create the checkerboard mask

def checkerboard_binary_mask(shape, orientation=0):
    height, width = shape[0], shape[1]
    height_range = tf.range(height)
    width_range = tf.range(width)
    height_odd_inx = tf.cast(tf.math.mod(height_range, 2), dtype=tf.bool)
    width_odd_inx = tf.cast(tf.math.mod(width_range, 2), dtype=tf.bool)
    odd_rows = tf.tile(tf.expand_dims(height_odd_inx, -1), [1, width])
    odd_cols = tf.tile(tf.expand_dims(width_odd_inx, 0), [height, 1])
    checkerboard_mask = tf.math.logical_xor(odd_rows, odd_cols)
    if orientation == 1:
        checkerboard_mask = tf.math.logical_not(checkerboard_mask)
    return tf.cast(tf.expand_dims(checkerboard_mask, -1), tf.float32)
```

This function creates a rank-3 Tensor to mask the `height`, `width` and `channels` dimensions of the input. We can take a look at this checkerboard mask for some example inputs below. In order to make the Tensors easier to inspect, we will squeeze out the single channel dimension (which is always 1 for this mask).

```python
# Run the checkerboard_binary_mask function to see an example
# NB: we squeeze the shape for easier viewing. The full shape is (4, 4, 1)

tf.squeeze(checkerboard_binary_mask((4, 4), orientation=0))
```

    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[0., 1., 0., 1.],
           [1., 0., 1., 0.],
           [0., 1., 0., 1.],
           [1., 0., 1., 0.]], dtype=float32)>

```python
# The `orientation` should be 0 or 1, and determines which way round the binary entries are

tf.squeeze(checkerboard_binary_mask((4, 4), orientation=1))
```

    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[1., 0., 1., 0.],
           [0., 1., 0., 1.],
           [1., 0., 1., 0.],
           [0., 1., 0., 1.]], dtype=float32)>

```python
checkerboard_binary_mask((4, 4), orientation=1)
```

    <tf.Tensor: shape=(4, 4, 1), dtype=float32, numpy=
    array([[[1.],
            [0.],
            [1.],
            [0.]],

           [[0.],
            [1.],
            [0.],
            [1.]],

           [[1.],
            [0.],
            [1.],
            [0.]],

           [[0.],
            [1.],
            [0.],
            [1.]]], dtype=float32)>

You should now complete the following function to create a channel-wise mask. This function takes a single integer `num_channels` as an input, as well as an `orientation` argument, similar to above. You can assume that the `num_channels` integer is even.

The function should return a rank-3 Tensor with singleton entries for `height` and `width`. In the channel axis, the first `num_channels // 2` entries should be zero (for `orientation=0`) and the final `num_channels // 2` entries should be one (for `orientation=0`). The zeros and ones should be reversed for `orientation=1`. The `dtype` of the returned Tensor should be `tf.float32`.

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def channel_binary_mask(num_channels, orientation=0):
    """
    This function takes an integer num_channels and orientation (0 or 1) as
    arguments. It should create a channel-wise binary mask with
    dtype=tf.float32, according to the above specification.
    The function should then return the binary mask.
    """
    # Should do tf.ones(num_channels-z.shape[0]) where z = tf.zeros(num_channels // 2) because
    # num_channels may not always be EVEN!

    z = tf.zeros(num_channels // 2)
    o = tf.ones(num_channels - z.shape[0])
    return tf.concat([z,o], axis = -1)[None, None, ...] if not orientation else tf.concat([o,z], axis = -1)[None, None,...]
```

```python
# Run your function to see an example channel-wise binary mask

channel_binary_mask(6, orientation=0)
channel_binary_mask(6, orientation=1)
channel_binary_mask(5, orientation=1)
```

    <tf.Tensor: shape=(1, 1, 5), dtype=float32, numpy=array([[[1., 1., 1., 0., 0.]]], dtype=float32)>

$$
\begin{align}
\text{Forward pass:}\qquad y &= b\odot x + (1-b)\odot\left(x\odot\exp(s(b\odot x)) + t(b\odot x)\right),\tag{5}\\
\text{Inverse pass:}\qquad x &= b\odot y + (1-b)\odot\left(y - t(b\odot x)) \odot\exp( -s(b\odot x)\right).\tag{6}
\end{align}
$$

```python
#### GRADED CELL ####

# Complete the following functions.
# Make sure to not change the function names or arguments.

def forward(x, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should implement the forward transformation in equation (5)
    and return the output Tensor y, which will have the same shape as x
    """
    shift, log_scale = shift_and_log_scale_fn(x * b)
    return x * b + (1 - b) * (x * tf.exp(log_scale) + shift)

def inverse(y, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should implement the forward transformation in equation (5)
    and return the output Tensor y, which will have the same shape as x
    """
    shift, log_scale = shift_and_log_scale_fn(y * b)
    return y * b + (1 - b) * (y - shift) * tf.exp(-log_scale)

```

```python
x = tf.random.normal((16, 32, 32, 3))
print(x.shape[-1])
b = channel_binary_mask(x.shape[-1], 1)
print(b)

forward(x, b, conv_resnet).shape
```

    3
    tf.Tensor([[[1. 1. 0.]]], shape=(1, 1, 3), dtype=float32)

    TensorShape([16, 32, 32, 3])

The new bijector class also requires the `log_det_jacobian` methods to be implemented. Recall that the log of the Jacobian determinant of the forward transformation is given by $\sum_{j}s(x_{1:d})_j$, where $s$ is the log-scale function of the affine coupling layer.

You should now complete the following functions to define the `forward_log_det_jacobian` and `inverse_log_det_jacobian` methods of the affine coupling layer bijector.

* Both functions `forward_log_det_jacobian` and `inverse_log_det_jacobian` takes an input Tensor `x` (or `y`), a rank-3 binary mask `b`, and the `shift_and_log_scale_fn` callable
* These arguments are the same as the description for the `forward` and `inverse` functions
* The `forward_log_det_jacobian` function should implement the log of the Jacobian determinant for the transformation $(5)$
* The `inverse_log_det_jacobian` function should implement the log of the Jacobian determinant for the transformation $(6)$
* Both functions should reduce sum over the last three axes of the input Tensor (`height`, `width` and `channels`)

```python
#### GRADED CELL ####

# Complete the following functions.
# Make sure to not change the function names or arguments.

def forward_log_det_jacobian(x, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should compute and return the log of the Jacobian determinant
    of the forward transformation in equation (5)
    """
    shift, log_scale = shift_and_log_scale_fn(x * b)
    return tf.reduce_sum(log_scale * (1-b), axis=[-1,-2,-3])

def inverse_log_det_jacobian(y, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor y, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should compute and return the log of the Jacobian determinant
    of the forward transformation in equation (6)
    """
    shift, log_scale = shift_and_log_scale_fn(y * b)
    return -tf.reduce_sum(log_scale * (1-b), axis=[-1,-2,-3])

```

You are now ready to create the coupling layer bijector, using bijector subclassing. You should complete the class below to define the `AffineCouplingLayer`.

* You should complete the initialiser `__init__`, and the internal class method `_get_mask`
* The `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods are completed for you using the functions you have written above. Do not modify these methods
* The initialiser takes the `shift_and_log_scale_fn` callable, `mask_type` string (either `"checkerboard"` or `"channel"`, `orientation` (integer, either `0` or `1`) as required arguments, and allows for extra keyword arguments
  * The required arguments should be set as class attributes in the initialiser (note that the `shift_and_log_scale_fn` attribute is being used in the `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods)
  * The initialiser should call the base class initialiser, and pass in any extra keyword arguments
  * The class should have a required number of event dimensions equal to 3
* The internal method `_get_mask` takes a `shape` as an argument, which is the shape of an input Tensor
  * This method should use the `checkerboard_binary_mask` and `channel_binary_mask` functions above, as well as the `mask_type` and `orientation` arguments passed to the initialiser to compute and return the required binary mask
  * This method is used in each of the `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods

```python
#### GRADED CELL ####

# Complete the following class.
# Make sure to not change the class or method names or arguments.

class AffineCouplingLayer(tfb.Bijector):
    """
    Class to implement the affine coupling layer.
    Complete the __init__ and _get_mask methods according to the instructions above.
    """

    def __init__(self, shift_and_log_scale_fn, mask_type, orientation, **kwargs):
        """
        The class initialiser takes the shift_and_log_scale_fn callable, mask_type,
        orientation and possibly extra keywords arguments. It should call the
        base class initialiser, passing any extra keyword arguments along.
        It should also set the required arguments as class attributes.
        """
        super().__init__(**kwargs, forward_min_event_ndims = 3)
        self.shift_and_log_scale_fn = shift_and_log_scale_fn
        self.mask_type = mask_type
        self.orientation = orientation

    def _get_mask(self, shape):
        """
        This internal method should use the binary mask functions above to compute
        and return the binary mask, according to the arguments passed in to the
        initialiser.
        """
        if self.mask_type == 'checkerboard':
            return checkerboard_binary_mask(shape[1:], self.orientation)
        else:
            return channel_binary_mask(shape[-1], self.orientation)

    def _forward(self, x):
        b = self._get_mask(x.shape)
        return forward(x, b, self.shift_and_log_scale_fn)

    def _inverse(self, y):
        b = self._get_mask(y.shape)
        return inverse(y, b, self.shift_and_log_scale_fn)

    def _forward_log_det_jacobian(self, x):
        b = self._get_mask(x.shape)
        return forward_log_det_jacobian(x, b, self.shift_and_log_scale_fn)

    def _inverse_log_det_jacobian(self, y):
        b = self._get_mask(y.shape)
        return inverse_log_det_jacobian(y, b, self.shift_and_log_scale_fn)
```

```python
# Test your function by creating an instance of the AffineCouplingLayer class

affine_coupling_layer = AffineCouplingLayer(conv_resnet, 'channel', orientation=1,
                                            name='affine_coupling_layer')
```

```python
# The following should return a Tensor of the same shape as the input

affine_coupling_layer.forward(tf.random.normal((16, 32, 32, 3))).shape
```

    TensorShape([16, 32, 32, 3])

```python
# The following should compute a log_det_jacobian for each event in the batch

affine_coupling_layer.forward_log_det_jacobian(tf.random.normal((16, 32, 32, 3)), event_ndims=3).shape
```

    TensorShape([16])

#### Combining the affine coupling layers

In the affine coupling layer, part of the input remains unchanged in the transformation $(5)$. In order to allow transformation of all of the input, several coupling layers are composed, with the orientation of the mask being reversed in subsequent layers.

Figure 2. RealNVP alternates the orientation of masks from one affine coupling layer to the next. From the original paper.

Our model design will be similar to the original architecture; we will compose three affine coupling layers with checkerboard masking, followed by a batch normalization bijector (`tfb.BatchNormalization` is a built-in bijector), followed by a squeezing operation, followed by three more affine coupling layers with channel-wise masking and a final batch normalization bijector.

The squeezing operation divides the spatial dimensions into 2x2 squares, and reshapes a Tensor of shape `(H, W, C)` into a Tensor of shape `(H // 2, W // 2, 4 * C)` as shown in Figure 1.

The squeezing operation is also a bijective operation, and has been provided for you in the class below.

```python
# Bijector class for the squeezing operation

class Squeeze(tfb.Bijector):

    def __init__(self, name='Squeeze', **kwargs):
        super(Squeeze, self).__init__(forward_min_event_ndims=3, is_constant_jacobian=True,
                                      name=name, **kwargs)

    def _forward(self, x):
        input_shape = x.shape
        height, width, channels = input_shape[-3:]
        y = tfb.Reshape((height // 2, 2, width // 2, 2, channels), event_shape_in=(height, width, channels))(x)
        y = tfb.Transpose(perm=[0, 2, 1, 3, 4])(y)
        y = tfb.Reshape((height // 2, width // 2, 4 * channels),
                        event_shape_in=(height // 2, width // 2, 2, 2, channels))(y)
        return y

    def _inverse(self, y):
        input_shape = y.shape
        height, width, channels = input_shape[-3:]
        x = tfb.Reshape((height, width, 2, 2, channels // 4), event_shape_in=(height, width, channels))(y)
        x = tfb.Transpose(perm=[0, 2, 1, 3, 4])(x)
        x = tfb.Reshape((2 * height, 2 * width, channels // 4),
                        event_shape_in=(height, 2, width, 2, channels // 4))(x)
        return x

    def _forward_log_det_jacobian(self, x):
        return tf.constant(0., x.dtype)

    def _inverse_log_det_jacobian(self, y):
        return tf.constant(0., y.dtype)

    def _forward_event_shape_tensor(self, input_shape):
        height, width, channels = input_shape[-3], input_shape[-2], input_shape[-1]
        return height // 2, width // 2, 4 * channels

    def _inverse_event_shape_tensor(self, output_shape):
        height, width, channels = output_shape[-3], output_shape[-2], output_shape[-1]
        return height * 2, width * 2, channels // 4
```

You can see the effect of the squeezing operation on some example inputs in the cells below. In the forward transformation, each spatial dimension is halved, whilst the channel dimension is multiplied by 4. The opposite happens in the inverse transformation.

```python
# Test the Squeeze bijector

squeeze = Squeeze()
squeeze(tf.ones((10, 32, 32, 3))).shape
```

    TensorShape([10, 16, 16, 12])

```python
# Test the inverse operation

squeeze.inverse(tf.ones((10, 4, 4, 96))).shape
```

    TensorShape([10, 8, 8, 24])

We can now construct a block of coupling layers according to the architecture described above. You should complete the following function to chain together the bijectors that we have constructed, to form a bijector that performs the following operations in the forward transformation:

* Three `AffineCouplingLayer` bijectors with `"checkerboard"` masking with orientations `0, 1, 0` respectively
* A `BatchNormalization` bijector
* A `Squeeze` bijector
* Three more `AffineCouplingLayer` bijectors with `"channel"` masking with orientations `0, 1, 0` respectively
* Another `BatchNormalization` bijector

The function takes the following arguments:
* `shift_and_log_scale_fns`: a list or tuple of six conv_resnet models
  * The first three models in this list are used in the three coupling layers with checkerboard masking
  * The last three models in this list are used in the three coupling layers with channel masking
* `squeeze`: an instance of the `Squeeze` bijector

_NB: at this point, we would like to point out that we are following the exposition in the original paper, and think of the forward transformation as acting on the input image. Note that this is in contrast to the convention of using the forward transformation for sampling, and the inverse transformation for computing log probs._

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def realnvp_block(shift_and_log_scale_fns, squeeze):
    """
    This function takes a list or tuple of six conv_resnet models, and an
    instance of the Squeeze bijector.
    The function should construct the chain of bijectors described above,
    using the conv_resnet models in the coupling layers.
    The function should then return the chained bijector.
    """
    # The question actually wants to stack the AffineCouplingLayer and Squeeze bijectors together
    # The inputs are the convet and squeeze functions which have NOT been converted to bijectors

    len_of_list = len(shift_and_log_scale_fns)
    middle = len_of_list // 2
    zero_one_zero = [0,1,0]
    lst1, lst2 = [], []

    for i,orientation in zip(range(middle),zero_one_zero):
        lst1.append(AffineCouplingLayer(shift_and_log_scale_fns[i],'checkerboard',
                   orientation = orientation))
    lst1.append(squeeze)
    lst1.append(tfb.BatchNormalization())
    for i,orientation in zip(range(middle,len_of_list),zero_one_zero):
        lst2.append(AffineCouplingLayer(shift_and_log_scale_fns[i],'checkerboard',
                   orientation = orientation))
    lst = lst1+lst2
    lst.append(tfb.BatchNormalization())
    return tfb.Chain(lst[::-1])
```

```python
# Run your function to create an instance of the bijector

checkerboard_fns = []
for _ in range(3):
    checkerboard_fns.append(get_conv_resnet((32, 32, 3), 512))
channel_fns = []
for _ in range(3):
    channel_fns.append(get_conv_resnet((16, 16, 12), 512))

block = realnvp_block(checkerboard_fns + channel_fns, squeeze)
```

```python
# Test the bijector on a dummy input

block.forward(tf.random.normal((10, 32, 32, 3))).shape
```

    TensorShape([10, 16, 16, 12])

#### Multiscale architecture

The final component of the RealNVP is the multiscale architecture. The squeeze operation reduces the spatial dimensions but increases the channel dimensions. After one of the blocks of coupling-squeeze-coupling that you have implemented above, half of the dimensions are factored out as latent variables, while the other half is further processed through subsequent layers. This results in latent variables that represent different scales of features in the model.

Figure 3. RealNVP creates latent variables at different scales by factoring out half of the dimensions at each scale. From the original paper.

The final scale does not use the squeezing operation, and instead applies four affine coupling layers with alternating checkerboard masks.

The multiscale architecture for two latent variable scales is implemented for you in the following bijector.

```python
# Bijector to implement the multiscale architecture

class RealNVPMultiScale(tfb.Bijector):

    def __init__(self, **kwargs):
        super(RealNVPMultiScale, self).__init__(forward_min_event_ndims=3, **kwargs)

        # First level
        shape1 = (32, 32, 3)  # Input shape
        shape2 = (16, 16, 12)  # Shape after the squeeze operation
        shape3 = (16, 16, 6)  # Shape after factoring out the latent variable
        self.conv_resnet1 = get_conv_resnet(shape1, 64)
        self.conv_resnet2 = get_conv_resnet(shape1, 64)
        self.conv_resnet3 = get_conv_resnet(shape1, 64)
        self.conv_resnet4 = get_conv_resnet(shape2, 128)
        self.conv_resnet5 = get_conv_resnet(shape2, 128)
        self.conv_resnet6 = get_conv_resnet(shape2, 128)
        self.squeeze = Squeeze()
        self.block1 = realnvp_block([self.conv_resnet1, self.conv_resnet2,
                                    self.conv_resnet3, self.conv_resnet4,
                                    self.conv_resnet5, self.conv_resnet6], self.squeeze)

        # Second level
        self.conv_resnet7 = get_conv_resnet(shape3, 128)
        self.conv_resnet8 = get_conv_resnet(shape3, 128)
        self.conv_resnet9 = get_conv_resnet(shape3, 128)
        self.conv_resnet10 = get_conv_resnet(shape3, 128)
        self.coupling_layer1 = AffineCouplingLayer(self.conv_resnet7, 'checkerboard', 0)
        self.coupling_layer2 = AffineCouplingLayer(self.conv_resnet8, 'checkerboard', 1)
        self.coupling_layer3 = AffineCouplingLayer(self.conv_resnet9, 'checkerboard', 0)
        self.coupling_layer4 = AffineCouplingLayer(self.conv_resnet10, 'checkerboard', 1)
        self.block2 = tfb.Chain([self.coupling_layer4, self.coupling_layer3,
                                 self.coupling_layer2, self.coupling_layer1])

    def _forward(self, x):
        h1 = self.block1.forward(x)
        z1, h2 = tf.split(h1, 2, axis=-1)
        z2 = self.block2.forward(h2)
        return tf.concat([z1, z2], axis=-1)

    def _inverse(self, y):
        z1, z2 = tf.split(y, 2, axis=-1)
        h2 = self.block2.inverse(z2)
        h1 = tf.concat([z1, h2], axis=-1)
        return self.block1.inverse(h1)

    def _forward_log_det_jacobian(self, x):
        log_det1 = self.block1.forward_log_det_jacobian(x, event_ndims=3)
        h1 = self.block1.forward(x)
        _, h2 = tf.split(h1, 2, axis=-1)
        log_det2 = self.block2.forward_log_det_jacobian(h2, event_ndims=3)
        return log_det1 + log_det2

    def _inverse_log_det_jacobian(self, y):
        z1, z2 = tf.split(y, 2, axis=-1)
        h2 = self.block2.inverse(z2)
        log_det2 = self.block2.inverse_log_det_jacobian(z2, event_ndims=3)
        h1 = tf.concat([z1, h2], axis=-1)
        log_det1 = self.block1.inverse_log_det_jacobian(h1, event_ndims=3)
        return log_det1 + log_det2

    def _forward_event_shape_tensor(self, input_shape):
        height, width, channels = input_shape[-3], input_shape[-2], input_shape[-1]
        return height // 4, width // 4, 16 * channels

    def _inverse_event_shape_tensor(self, output_shape):
        height, width, channels = output_shape[-3], output_shape[-2], output_shape[-1]
        return 4 * height, 4 * width, channels // 16
```

```python
# Create an instance of the multiscale architecture

multiscale_bijector = RealNVPMultiScale()
```

#### Data preprocessing bijector

We will also preprocess the image data before sending it through the RealNVP model. To do this, for a Tensor $x$ of pixel values in $[0, 1]^D$, we transform $x$ according to the following:

$$
T(x) = \text{logit}\left(\alpha + (1 - 2\alpha)x\right),\tag{7}
$$

where $\alpha$ is a parameter, and the logit function is the inverse of the sigmoid function, and is given by

$$
\text{logit}(p) = \log (p) - \log (1 - p).
$$

You should now complete the following function to construct this bijector from in-built bijectors from the bijectors module.

* The function takes the parameter `alpha` as an input, which you can assume to take a small positive value ($\ll0.5$)
* The function should construct and return a bijector that computes $(7)$ in the forward pass

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_preprocess_bijector(alpha):
    """
    This function should create a chained bijector that computes the
    transformation T in equation (7) above.
    This can be computed using in-built bijectors from the bijectors module.
    Your function should then return the chained bijector.
    """

    return tfb.Chain([tfb.Invert(tfb.Sigmoid()),
               tfb.Shift(shift=alpha),
               tfb.Scale(scale=(1-2*alpha))
    ])
```

```python
# Create an instance of the preprocess bijector

preprocess = get_preprocess_bijector(0.05)
```

#### Train the RealNVP model

Finally, we will use our RealNVP model to train

We will use the following model class to help with the training process.

```python
# Helper class for training

class RealNVPModel(Model):

    def __init__(self, **kwargs):
        super(RealNVPModel, self).__init__(**kwargs)
        self.preprocess = get_preprocess_bijector(0.05)
        self.realnvp_multiscale = RealNVPMultiScale()
        self.bijector = tfb.Chain([self.realnvp_multiscale, self.preprocess])

    def build(self, input_shape):
        output_shape = self.bijector(tf.expand_dims(tf.zeros(input_shape[1:]), axis=0)).shape
        self.base = tfd.Independent(tfd.Normal(loc=tf.zeros(output_shape[1:]), scale=1.),
                                    reinterpreted_batch_ndims=3)
        self._bijector_variables = (
            list(self.bijector.variables))
        self.flow = tfd.TransformedDistribution(
            distribution=self.base,
            bijector=tfb.Invert(self.bijector),
        )
        super(RealNVPModel, self).build(input_shape)

    def call(self, inputs, training=None, **kwargs):
        return self.flow

    def sample(self, batch_size):
        sample = self.base.sample(batch_size)
        return self.bijector.inverse(sample)
```

```python
# Create an instance of the RealNVPModel class

realnvp_model = RealNVPModel()
realnvp_model.build((1, 32, 32, 3))
```

```python
# Compute the number of variables in the model

print("Total trainable variables:")
print(sum([np.prod(v.shape) for v in realnvp_model.trainable_variables]))
```

    Total trainable variables:
    315198

Note that the model's `call` method returns the `TransformedDistribution` object. Also, we have set up our datasets to return the input image twice as a 2-tuple. This is so we can train our model with negative log-likelihood as normal.

```python
# Define the negative log-likelihood loss function

def nll(y_true, y_pred):
    return -y_pred.log_prob(y_true)
```

It is recommended to use the GPU accelerator hardware on Colab to train this model, as it can take some time to train. Note that it is not required to train the model in order to pass this assignment. For optimal results, a larger model should be trained for longer.

```python
# Compile and train the model

realnvp_model.compile(loss=nll, optimizer=Adam())
realnvp_model.fit(train_ds, validation_data=val_ds, epochs=20)
```

```python
# Evaluate the model

realnvp_model.evaluate(test_ds)
```

#### Generate some samples

```python
# Sample from the model

samples = realnvp_model.sample(8).numpy()
```

```python
# Display the samples

n_img = 8
f, axs = plt.subplots(2, n_img // 2, figsize=(14, 7))

for k, image in enumerate(samples):
    i = k % 2
    j = k // 2
    axs[i, j].imshow(np.clip(image, 0., 1.))
    axs[i, j].axis('off')
f.subplots_adjust(wspace=0.01, hspace=0.03)
```

Congratulations on completing this programming assignment! In the next week of the course we will look at the variational autoencoder.

Probabilistic Deep Learning with TensorFlow 2/.ipynb_checkpoints/Week 3 Programming Assignment-Completed-checkpoint.ipynb


# Programming Assignment

## VAE for the CelebA dataset

### Instructions
In this programming assignment, you will implement the variational autoencoder algorithm for an image dataset of celebrity faces. You will use the trained encoder and decoder networks to reconstruct and generate images. You will also see how the latent space encodes high-level information about the images.

Some code cells are provided for you in the notebook. You should avoid editing provided code, and make sure to execute the cells in order to avoid unexpected errors. Some cells begin with the line:

`#### GRADED CELL ####`

Don't move or edit this first line - this is what the automatic grader looks for to recognise graded cells. These cells require you to write your own code to complete them, and are automatically graded when you submit the notebook. Don't edit the function name or signature provided in these cells, otherwise the automatic grader might not function properly.

### How to submit

Complete all the tasks you are asked for in the worksheet. When you have finished and are happy with your code, press the Submit Assignment button at the top of this notebook.

### Let's get started!

We'll start running some imports, and loading the dataset. Do not edit the existing imports in the following cell. If you would like to make further Tensorflow imports, you should add them here.

```python
#### PACKAGE IMPORTS ####

# Run this cell first to import all required packages. Do not make any imports elsewhere in the notebook

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os

from tensorflow.keras import Sequential, Model
from tensorflow.keras.layers import (Dense, Flatten, Reshape, Concatenate, Conv2D,
                                     UpSampling2D, BatchNormalization)
tfd = tfp.distributions
tfb = tfp.bijectors
tfpl = tfp.layers

# If you would like to make further imports from tensorflow, add them here

```

#### The Large-scale CelebFaces Attributes (CelebA) Dataset

For this assignment you will use a subset of the CelebFaces Attributes (CelebA) dataset. The full dataset contains over 200K images CelebA contains thousands of colour images of the faces of celebrities, together with tagged attributes such as 'Smiling', 'Wearing glasses', or 'Wearing lipstick'. It also contains information about bounding boxes and facial part localisation. CelebA is a popular dataset that is commonly used for face attribute recognition, face detection, landmark (or facial part) localization, and face editing & synthesis.

* Z. Liu, P. Luo, X. Wang, and X. Tang. "Deep Learning Face Attributes in the Wild", Proceedings of International Conference on Computer Vision (ICCV), 2015.

You can read about the dataset in more detail [here](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html).

Your goal is to implement the variational autoencoder algorithm for a subset of the CelebA dataset. For practical reasons we will keep the dataset and the network size relatively small.

#### Import the data

The dataset required for this project can be downloaded from the following link:

https://drive.google.com/file/d/1qzbcYRjxO6yU_usb4CXyGObUEjIVYWAp/view?usp=sharing

You should unzip this file and store in Drive for use in this Colab notebook.

```python
# Run this cell to connect to your Drive folder

from google.colab import drive
drive.mount('/content/gdrive')
```

#### Load the dataset

The following functions will be useful for loading and preprocessing the dataset. The subset you will use for this assignment consists of 10,000 training images, 1000 validation images and 1000 test images. These examples have been chosen to respect the original training / validation / test split of the dataset.

```python
# Functions for loading and preprocessing the images

def load_image(filepath):
    raw_img = tf.io.read_file(filepath)
    img_tensor_int = tf.image.decode_jpeg(raw_img, channels=3)
    img_tensor_flt = tf.image.convert_image_dtype(img_tensor_int, tf.float32)
    return img_tensor_flt, img_tensor_flt

def load_dataset(split):
    train_list_ds = tf.data.Dataset.list_files('/path/to/data/{}/*.jpg'.format(split), shuffle=False)
    train_ds = train_list_ds.map(load_image)
    return train_ds
```

```python
# Load the training, validation and testing datasets splits

train_ds = load_dataset('train')
val_ds = load_dataset('val')
test_ds = load_dataset('test')
```

```python
# Display a few examples

n_examples_shown = 6
f, axs = plt.subplots(1, n_examples_shown, figsize=(16, 3))

for j, image in enumerate(train_ds.take(n_examples_shown)):
    axs[j].imshow(image[0])
    axs[j].axis('off')
```

```python
# Batch the Dataset objects

batch_size = 32
train_ds = train_ds.batch(batch_size)
val_ds = val_ds.batch(batch_size)
test_ds = test_ds.batch(batch_size)
```

#### Mixture of Gaussians distribution

We will define a prior distribution that is a mixture of Gaussians. This is a more flexible distribution that is comprised of $K$ separate Gaussians, that are combined together with some weighting assigned to each.

Recall that the probability density function for a multivariate Gaussian distribution with mean $\mu\in\mathbb{R}^D$ and covariance matrix $\Sigma\in\mathbb{R}^{D\times D}$ is given by

$$
\mathcal{N}(\mathbf{z}; \mathbf{\mu}, \Sigma) = \frac{1}{(2\pi)^{D/2}|\Sigma|^{1/2}}
\exp\left(-\frac{1}{2}(\mathbf{z}-\mathbf{\mu})^T\Sigma^{-1}(\mathbf{z}-\mathbf{\mu})\right).
$$

A mixture of Gaussians with $K$ components defines $K$ Gaussians defined by means $\mathbf{\mu}_k$ and covariance matrices $\Sigma_k$, for $k=1,\ldots,K$. It also requires mixing coefficients $\pi_k$, $k=1,\ldots,K$ with $\sum_{k} \pi_k = 1$. These coefficients define a categorical distribution over the $K$ Gaussian components. To sample an event, we first sample from the categorical distribution, and then again from the corresponding Gaussian component.

The probability density function of the mixture of Gaussians is simply the weighted sum of probability density functions for each Gaussian component:

$$
p(\mathbf{z}) = \sum_{k=1}^K \pi_k \mathcal{N}(\mathbf{z}; \mathbf{\mu}_k, \Sigma_k)
$$

#### Define the prior distribution

You should now complete the following function to define the mixture of Gaussians distribution for the prior, for a given number of components and latent space dimension. Each Gaussian component will have a diagonal covariance matrix. This distribution will have fixed mixing coefficients, but trainable means and standard deviations.

* The function takes `num_modes` (number of components) and `latent_dim` as arguments
* Use the `tfd.MixtureSameFamily` for the prior distribution. Take a look at [the documentation](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/MixtureSameFamily) for this distribution
  * The constructor takes a `mixture_distribution` and `components_distribution` as required arguments
  * The `mixture_distribution` should be fixed to a uniform `tfd.Categorical` distribution, so that $pi_k = 1/K$ in the above equation. This argument will therefore not contain any trainable variables
  * The `components_distribution` should be a `tfd.MultivariateNormalDiag` distribution batch shape equal to `[num_modes]` and event shape equal to `[latent_dim]`.
    * The `tfd.MultivariateNormalDiag` distribution should have trainable `loc` parameter (initialised with a random normal distribution)  and trainable `scale_diag` parameter (initialised to ones)
    * The `scale_diag` variable should be enforced to be positive using `tfp.util.TransformedVariable` and the `tfb.Softplus` bijection

The function should return the instance of the `tfd.MixtureSameFamily` distribution.

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_prior(num_modes, latent_dim):
    """
    This function should create an instance of a MixtureSameFamily distribution
    according to the above specification.
    The function takes the num_modes and latent_dim as arguments, which should
    be used to define the distribution.
    Your function should then return the distribution instance.
    """

```

```python
# Run your function to get the prior distribution with 2 components and latent_dim = 50

prior = get_prior(num_modes=2, latent_dim=50)
```

#### Define the encoder network

We will now define the encoder network as part of the VAE. First, we will define the `KLDivergenceRegularizer` to use in the encoder network to add the KL divergence part of the loss. This should be defined according to the following specification:

* The function takes the `prior_distribution` as an argument
* The function should use the `tfpl.KLDivergenceRegularizer` object to add the KL loss term
* The `tfpl.KLDivergenceRegularizer` should use a weight factor of 1.0 for the KL loss (standard ELBO objective)
* The KL loss cannot be computed exactly, so the `tfpl.KLDivergenceRegularizer` should compute a Monte Carlo approximation by drawing 3 samples from the posterior, and then averaging over the sample and batch axes

Your function should then return the `tfpl.KLDivergenceRegularizer` object.

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_kl_regularizer(prior_distribution):
    """
    This function should create an instance of the KLDivergenceRegularizer
    according to the above specification.
    The function takes the prior_distribution, which should be used to define
    the distribution.
    Your function should then return the KLDivergenceRegularizer instance.
    """

```

```python
# Run your function to get the KLDivergenceRegularizer

kl_regularizer = get_kl_regularizer(prior)
```

You should now complete the following function to define the encoder network, according to the following specification:

* The function takes the `latent_dim` and `kl_regularizer` as arguments
* Use the `Sequential` class to define the model
  * The first layer should be a Conv2D layer with 32 filters, 4x4 kernel size, ReLU activation, stride of 2x2, and 'SAME' padding. It should also set the `input_shape` to `(64, 64, 3)`
  * The second layer should be a BatchNormalization layer
  * The third layer should be a Conv2D layer with 64 filters, 4x4 kernel size, ReLU activation, stride of 2x2, and 'SAME' padding
  * The fourth layer should be a BatchNormalization layer
  * The fifth layer should be a Conv2D layer with 128 filters, 4x4 kernel size, ReLU activation, stride of 2x2, and 'SAME' padding
  * The sixth layer should be a BatchNormalization layer
  * The seventh layer should be a Conv2D layer with 256 filters, 4x4 kernel size, ReLU activation, stride of 2x2, and 'SAME' padding
  * The eighth layer should be a BatchNormalization layer
  * The ninth layer should be a Flatten layer
  * The tenth layer should be a Dense layer with no activation function, and the right number of units to parameterise a `MultivariateNormalTriL` layer with event size equal to `latent_dim`
  * The final layer should be a `MultivariateNormalTriL` layer with event size equal to `latent_dim`, and it should use the `kl_regularizer` passed in as the argument as the activity regularizer

In total, your model should have 11 layers.

The function should then return the encoder model.

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_encoder(latent_dim, kl_regularizer):
    """
    This function should build a CNN encoder model according to the above specification.
    The function takes latent_dim and kl_regularizer as arguments, which should be
    used to define the model.
    Your function should return the encoder model.
    """

```

```python
# Run your function to get the encoder

encoder = get_encoder(latent_dim=50, kl_regularizer=kl_regularizer)
```

```python
# Print the encoder summary

encoder.summary()
```

#### Define the decoder network

You should now define the decoder network for the VAE, using the `Sequential` API. This should be a neural network that returns an IndependentBernoulli distribution of `event_shape=(64, 64, 3)`.

* The function takes the `latent_dim` as an argument
* Use the `Sequential` class to define the model
  * The first layer should be a Dense layer with 4096 units and ReLU activation. It should also set the `input_shape` to `(latent_dim,)`
  * The second layer should be a Reshape layer, that reshapes its input to `(4, 4, 256)`
  * The third layer should be an UpSampling2D layer with upsampling factor of `(2, 2)`
  * The fourth layer should be a Conv2D layer with 128 filters, 3x3 kernel size, ReLU activation, and 'SAME' padding
  * The fifth layer should be an UpSampling2D layer with upsampling factor of `(2, 2)`
  * The sixth layer should be a Conv2D layer with 64 filters, 3x3 kernel size, ReLU activation, and 'SAME' padding
  * The seventh layer should be an UpSampling2D layer with upsampling factor of `(2, 2)`
  * The eighth layer should be a Conv2D layer with 32 filters, 3x3 kernel size, ReLU activation, and 'SAME' padding
  * The ninth layer should be an UpSampling2D layer with upsampling factor of `(2, 2)`
  * The tenth layer should be a Conv2D layer with 128 filters, 3x3 kernel size, ReLU activation, and 'SAME' padding
  * The eleventh layer should be a Conv2D layer with 3 filters, 3x3 kernel size, no activation function, and 'SAME' padding
  * The twelfth layer should be a Flatten layer
  * The final layer should be a `IndependentBernoulli` layer with event size equal to `(64, 64, 3)`

In total, your model should have 13 layers.

The function should then return the decoder model.

```python
 #### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_decoder(latent_dim):
    """
    This function should build a CNN decoder model according to the above specification.
    The function takes latent_dim as an argument, which should be used to define the model.
    Your function should return the decoder model.
    """

```

```python
# Run your function to get the decoder

decoder = get_decoder(latent_dim=50)
```

```python
# Print the decoder summary

decoder.summary()
```

#### Link the encoder and decoder together

The following cell connects `encoder` and `decoder` to form the end-to-end architecture.

```python
# Connect the encoder and decoder

vae = Model(inputs=encoder.inputs, outputs=decoder(encoder.outputs))
```

#### Define the average reconstruction loss

You should now define the reconstruction loss that forms the remaining part of the negative ELBO objective. This function should take a batch of images of shape `(batch_size, 64, 64, 3)` in the first argument, and the output of the decoder after passing the batch of images through `vae` in the second argument.

The loss should be defined so that it returns
\begin{equation}
    -\frac{1}{n}\sum_{i=1}^n \log p(x_i|z_i)
\end{equation}
where $n$ is the batch size and $z_i$ is sampled from $q(z|x_i)$, the encoding distribution a.k.a. the approximate posterior. The value of this expression is always a scalar.

Expression (1) is, as you know, is an estimate of the (negative of the) batch's average expected reconstruction loss:

\begin{equation}
    -\frac{1}{n}\sum_{i=1}^n \mathrm{E}_{Z\sim q(z|x_i)}\big[\log p(x_i|Z)\big]
\end{equation}

_Hints_ :
- _You may find the `log_prob` method of the decoding distribution helpful._
- _The code you add does not need to be more than one line._

```python
 #### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def reconstruction_loss(batch_of_images, decoding_dist):
    """
    This function should compute and return the average expected reconstruction loss,
    as defined above.
    The function takes batch_of_images (Tensor containing a batch of input images to
    the encoder) and decoding_dist (output distribution of decoder after passing the
    image batch through the encoder and decoder) as arguments.
    The function should return the scalar average expected reconstruction loss.
    """

```

#### Compile and fit the model

It's now time to compile and train the model. This may take some time, so as an alternative, you can also load a pre-trained model below if you wish.

The pre-trained model can be downloaded from the following link:

https://drive.google.com/file/d/15JMM4Pw4lnYvgR_5WOjW8ttkhglTGOGq/view?usp=sharing

You should unzip this file and store in Drive for use in this Colab notebook.

```python
# Compile the model

optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)
vae.compile(optimizer=optimizer, loss=reconstruction_loss)
```

```python
# EITHER... compile and fit the model

vae.fit(train_ds, validation_data=val_ds, epochs=20)
```

```python
# OR... load the pre-trained model

ckpt = tf.train.Checkpoint(model=vae)
ckpt.restore(tf.train.latest_checkpoint('/path/to/model'))
```

```python
# Evaluate the model on the test set

test_loss = vae.evaluate(test_ds)
print("Test loss: {}".format(test_loss))
```

#### Compute reconstructions of test images

We will now take a look at some image reconstructions from the encoder-decoder architecture.

You should complete the following function, that uses `encoder` and `decoder` to reconstruct images from the test dataset. This function takes the encoder, decoder and a Tensor batch of test images as arguments. The function should be completed according to the following specification:

* Get the mean of the encoding distributions from passing the batch of images into the encoder
* Pass these latent vectors through the decoder to get the output distribution

Your function should then return the mean of the output distribution, which will be a Tensor of shape `(batch_size, 64, 64, 3)`.

```python
 #### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def reconstruct(encoder, decoder, batch_of_images):
    """
    This function should compute reconstructions of the batch_of_images according
    to the above instructions.
    The function takes the encoder, decoder and batch_of_images as inputs, which
    should be used to compute the reconstructions.
    The function should then return the reconstructions Tensor.
    """

```

```python
# Run your function to compute reconstructions of random samples from the test dataset

n_reconstructions = 7
files = [f for f in os.listdir('./data/test/') if f.endswith('.jpg')]
num_test_files = len(files)
test_ds_for_reconstructions = load_dataset('test')
for all_test_images, _ in test_ds_for_reconstructions.batch(num_test_files).take(1):
    all_test_images_np = all_test_images.numpy()
example_images = all_test_images_np[np.random.choice(num_test_files, n_reconstructions, replace=False)]

reconstructions = reconstruct(encoder, decoder, example_images).numpy()
```

```python
# Plot the reconstructions

f, axs = plt.subplots(2, n_reconstructions, figsize=(16, 6))
axs[0, n_reconstructions // 2].set_title("Original test images")
axs[1, n_reconstructions // 2].set_title("Reconstructed images")
for j in range(n_reconstructions):
    axs[0, j].imshow(example_images[j])
    axs[1, j].imshow(reconstructions[j])
    axs[0, j].axis('off')
    axs[1, j].axis('off')

plt.tight_layout();
```

#### Sample new images from the generative model

Now we will sample from the generative model; that is, first sample latent vectors from the prior, and then decode those latent vectors with the decoder.

You should complete the following function to generate new images. This function takes the prior distribution and decoder network as arguments, as well as the number of samples to generate. This function should be completed according to the following:

* Sample a batch of `n_samples` images from the prior distribution, to obtain a latent vector Tensor of shape `(n_samples, 50)`
* Pass this batch of latent vectors through the decoder, to obtain an Independent Bernoulli distribution with batch shape equal to `[n_samples]` and event shape equal to `[64, 64, 3]`.

The function should then return the mean of the Bernoulli distribution, which will be a Tensor of shape `(n_samples, 64, 64, 3)`.

```python
 #### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def generate_images(prior, decoder, n_samples):
    """
    This function should compute generate new samples of images from the generative model,
    according to the above instructions.
    The function takes the prior distribution, decoder and number of samples as inputs, which
    should be used to generate the images.
    The function should then return the batch of generated images.
    """

```

```python
# Run your function to generate new images

n_samples = 10
sampled_images = generate_images(prior, decoder, n_samples)

f, axs = plt.subplots(1, n_samples, figsize=(16, 6))

for j in range(n_samples):
    axs[j].imshow(sampled_images[j])
    axs[j].axis('off')

plt.tight_layout();
```

### Modify generations with attribute vector

This final section is ungraded, but you will hopefully find it to be an interesting extension. We will see how the latent space encodes high-level information about the images, even though it has not been trained with any information apart from the images themselves.

As mentioned in the introduction, each image in the CelebA dataset is labelled according to the attributes of the person pictured.

Run the cells below to load these labels, along with a subset of the training data.

```python
# Function to load labels and images as a numpy array

def load_labels_and_image_arrays(split):
    dataset = load_dataset(split)
    files = [f for f in os.listdir('/path/to/data/{}/'.format(split)) if f.endswith('.jpg')]

    for all_images, _ in dataset.batch(len(files)).take(1):
        all_images_np = all_images.numpy()

    labels = pd.read_csv('/path/to/list_attr_celeba_subset.csv')
    labels = labels[labels['image_id'].isin(files)]
    return labels, all_images_np
```

```python
# Load labels in a pandas DataFrame, training_subset is a numpy array

train_labels, training_subset = load_labels_and_image_arrays('train')
```

```python
# List the attributes contained in the DataFrame

train_labels.columns[2:]
```

Each image is labelled with a binary indicator (1 is True, -1 is False) according to whether it possesses the attribute.

```python
# View a sample from the labels data

train_labels.sample(5)
```

```python
# Select an attribute

attribute = 'Smiling'
```

```python
# Separate the images into those that have the attribute, and those that don't

attribute_mask = (train_labels[attribute] == 1)
images_with_attribute = training_subset[attribute_mask]

not_attribute_mask = (train_labels[attribute] == -1)
images_without_attribute = training_subset[not_attribute_mask]
```

#### Get the 'attribute vector'

We will now encode each of the images that have the chosen attribute into the latent space by passing them through the encoder. We then average the latent codes obtained for all of these images to obtain a single latent code.

We then do the same for the images that do not have the chosen attribute. This gives an average latent code for images with the attribute, and an average latent code for images without the attribute. Intuitively speaking, the difference between these two vectors then gives us the 'direction' in latent space that corresponds to the presence of the attribute.

```python
# Encode the images with and without the chosen attribute

encoded_images_with_attribute = encoder(images_with_attribute)
encoded_images_without_attribute = encoder(images_without_attribute)
```

```python
# Average the latent vectors for each batch of encodings

mean_encoded_images_with_attribute = tf.reduce_mean(encoded_images_with_attribute.mean(),
                                                    axis=0, keepdims=True)
mean_encoded_images_without_attribute = tf.reduce_mean(encoded_images_without_attribute.mean(),
                                                    axis=0, keepdims=True)
```

```python
# Get the attribute vector

attribute_vector = mean_encoded_images_with_attribute -  mean_encoded_images_without_attribute
```

We can view this attribute vector by decoding it:

```python
# Display the decoded attribute vector

decoded_a = decoder(attribute_vector).mean()
plt.imshow(decoded_a.numpy().squeeze())
plt.axis('off');
```

#### Modify reconstructions using the attribute vector

We can now use the attribute vector to add the attribute to an image reconstruction, where that attribute wasn't present before. To do this, we can just add the attribute vector to the latent vector encoding of the image, and then decode the result. We can also adjust the strength of the attribute vector by scaling with a multiplicative parameter.

```python
# Add the attribute vector to a sample of images that don't have the attribute

n_examples = 7
sampled_inx = np.random.choice(images_without_attribute.shape[0], n_examples, replace=False)
sample_images_without_attribute = images_without_attribute[sampled_inx]
sample_images_encodings = encoder(sample_images_without_attribute)
sample_images_reconstructions = decoder(sample_images_encodings).mean()

k = 2.5  # Weighting of attribute vector
modified_sample_images_encodings = sample_images_encodings + (k * attribute_vector)
modified_reconstructions = decoder(modified_sample_images_encodings).mean()
```

```python
# Display the original images, their reconstructions, and modified reconstructions

f, axs = plt.subplots(3, n_examples, figsize=(16, 6))
axs[0, n_examples // 2].set_title("Original images")
axs[1, n_examples // 2].set_title("Reconstructed images")
axs[2, n_examples // 2].set_title("Images with added attribute")
for j in range(n_examples):
    axs[0, j].imshow(sample_images_without_attribute[j])
    axs[1, j].imshow(sample_images_reconstructions[j])
    axs[2, j].imshow(modified_reconstructions[j])
    for ax in axs[:, j]: ax.axis('off')

plt.tight_layout();
```

You could also try removing the attribute from images that possess the attribute, or experiment with a different attribute.

---

Congratulations for completing this programming assignment! You're now ready to move on to the capstone project for this course.

ICL_Probabilistic_Deep_Learning_with_TensorFlow/Week 4 Variational Autoencoder.ipynb