# Homework and bake-off: word-level entailment with neural networks

```python
__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2020"
```

## Contents

1. [Overview](#Overview)
1. [Set-up](#Set-up)
1. [Data](#Data)
  1. [Edge disjoint](#Edge-disjoint)
  1. [Word disjoint](#Word-disjoint)
1. [Baseline](#Baseline)
  1. [Representing words: vector_func](#Representing-words:-vector_func)
  1. [Combining words into inputs: vector_combo_func](#Combining-words-into-inputs:-vector_combo_func)
  1. [Classifier model](#Classifier-model)
  1. [Baseline results](#Baseline-results)
1. [Homework questions](#Homework-questions)
  1. [Hypothesis-only baseline [2 points]](#Hypothesis-only-baseline-[2-points])
  1. [Alternatives to concatenation [2 points]](#Alternatives-to-concatenation-[2-points])
  1. [A deeper network [2 points]](#A-deeper-network-[2-points])
  1. [Your original system [3 points]](#Your-original-system-[3-points])
1. [Bake-off [1 point]](#Bake-off-[1-point])

## Overview

The general problem is word-level natural language inference.

Training examples are pairs of words $(w_{L}, w_{R}), y$ with $y = 1$ if $w_{L}$ entails $w_{R}$, otherwise $0$.

The homework questions below ask you to define baseline models for this and develop your own system for entry in the bake-off, which will take place on a held-out test-set distributed at the start of the bake-off. (Thus, all the data you have available for development is available for training your final system before the bake-off begins.)

## Set-up

See [the first notebook in this unit](nli_01_task_and_data.ipynb) for set-up instructions.

```python
from collections import defaultdict
import json
import numpy as np
import os
import pandas as pd
from torch_shallow_neural_classifier import TorchShallowNeuralClassifier
import nli
import utils
import io
```

```python
DATA_HOME = 'data'

NLIDATA_HOME = os.path.join(DATA_HOME, 'nlidata')

wordentail_filename = os.path.join(
    NLIDATA_HOME, 'nli_wordentail_bakeoff_data.json')

GLOVE_HOME = os.path.join(DATA_HOME, 'glove.6B')
```

## Data

I've processed the data into two different train/test splits, in an effort to put some pressure on our models to actually learn these semantic relations, as opposed to exploiting regularities in the sample.

* `edge_disjoint`: The `train` and `dev` __edge__ sets are disjoint, but many __words__ appear in both `train` and `dev`.
* `word_disjoint`: The `train` and `dev` __vocabularies are disjoint__, and thus the edges are disjoint as well.

These are very different problems. For `word_disjoint`, there is real pressure on the model to learn abstract relationships, as opposed to memorizing properties of individual words.

```python
with open(wordentail_filename) as f:
    wordentail_data = json.load(f)
```

The outer keys are the  splits plus a list giving the vocabulary for the entire dataset:

```python
wordentail_data.keys()
```

    dict_keys(['edge_disjoint', 'vocab', 'word_disjoint'])

### Edge disjoint

```python
wordentail_data['edge_disjoint'].keys()
```

    dict_keys(['dev', 'train'])

This is what the split looks like; all three have this same format:

```python
wordentail_data['edge_disjoint']['dev'][: 5]
```

    [[['sweater', 'stroke'], 0],
     [['constipation', 'hypovolemia'], 0],
     [['disease', 'inflammation'], 0],
     [['herring', 'animal'], 1],
     [['cauliflower', 'outlook'], 0]]

Let's test to make sure no edges are shared between `train` and `dev`:

```python
nli.get_edge_overlap_size(wordentail_data, 'edge_disjoint')
```

    0

As we expect, a *lot* of vocabulary items are shared between `train` and `dev`:

```python
nli.get_vocab_overlap_size(wordentail_data, 'edge_disjoint')
```

    2916

This is a large percentage of the entire vocab:

```python
len(wordentail_data['vocab'])
```

    8470

Here's the distribution of labels in the `train` set. It's highly imbalanced, which will pose a challenge for learning. (I'll go ahead and reveal that the `dev` set is similarly distributed.)

```python
def label_distribution(split):
    return pd.DataFrame(wordentail_data[split]['train'])[1].value_counts()
```

```python
label_distribution('edge_disjoint')
```

    0    14650
    1     2745
    Name: 1, dtype: int64

### Word disjoint

```python
wordentail_data['word_disjoint'].keys()
```

    dict_keys(['dev', 'train'])

In the `word_disjoint` split, no __words__ are shared between `train` and `dev`:

```python
nli.get_vocab_overlap_size(wordentail_data, 'word_disjoint')
```

    0

Because no words are shared between `train` and `dev`, no edges are either:

```python
nli.get_edge_overlap_size(wordentail_data, 'word_disjoint')
```

    0

The label distribution is similar to that of `edge_disjoint`, though the overall number of examples is a bit smaller:

```python
label_distribution('word_disjoint')
```

    0    7199
    1    1349
    Name: 1, dtype: int64

## Baseline

Even in deep learning, __feature representation is vital and requires care!__ For our task, feature representation has two parts: representing the individual words and combining those representations into a single network input.

### Representing words: vector_func

Let's consider two baseline word representations methods:

1. Random vectors (as returned by `utils.randvec`).
1. 50-dimensional GloVe representations.

```python
def randvec(w, n=50, lower=-1.0, upper=1.0):
    """Returns a random vector of length `n`. `w` is ignored."""
    return utils.randvec(n=n, lower=lower, upper=upper)
```

```python
# Any of the files in glove.6B will work here:

glove_dim = 50

glove_src = os.path.join(GLOVE_HOME, 'glove.6B.{}d.txt'.format(glove_dim))

# Creates a dict mapping strings (words) to GloVe vectors:
GLOVE = utils.glove2dict(glove_src)

def glove_vec(w):
    """Return `w`'s GloVe representation if available, else return
    a random vector."""
    #if w not in GLOVE:
    #    print(w)
    return GLOVE.get(w, randvec(w, n=glove_dim))
```

### Combining words into inputs: vector_combo_func

Here we decide how to combine the two word vectors into a single representation. In more detail, where `u` is a vector representation of the left word and `v` is a vector representation of the right word, we need a function `vector_combo_func` such that `vector_combo_func(u, v)` returns a new input vector `z` of dimension `m`. A simple example is concatenation:

```python
def vec_concatenate(u, v):
    """Concatenate np.array instances `u` and `v` into a new np.array"""
    return np.concatenate((u, v))
```

`vector_combo_func` could instead be vector average, vector difference, etc. (even combinations of those) – there's lots of space for experimentation here; [homework question 2](#Alternatives-to-concatenation-[1-point]) below pushes you to do some exploration.

### Classifier model

For a baseline model, I chose `TorchShallowNeuralClassifier`:

```python
net = TorchShallowNeuralClassifier(hidden_dim=50, max_iter=100)
```

### Baseline results

The following puts the above pieces together, using `vector_func=glove_vec`, since `vector_func=randvec` seems so hopelessly misguided for `word_disjoint`!

```python
word_disjoint_experiment = nli.wordentail_experiment(
    train_data=wordentail_data['word_disjoint']['train'],
    assess_data=wordentail_data['word_disjoint']['dev'],
    model=net,
    vector_func=glove_vec,
    vector_combo_func=vec_concatenate)
```

    Finished epoch 100 of 100; error is 0.02446776209399104

                  precision    recall  f1-score   support

               0      0.930     0.934     0.932      1910
               1      0.455     0.439     0.447       239

        accuracy                          0.879      2149
       macro avg      0.692     0.687     0.689      2149
    weighted avg      0.877     0.879     0.878      2149

```python
GLOVE.get("", "")
```

    ''

## Homework questions

Please embed your homework responses in this notebook, and do not delete any cells from the notebook. (You are free to add as many cells as you like as part of your responses.)

### Hypothesis-only baseline [2 points]

During our discussion of SNLI and MultiNLI, we noted that a number of research teams have shown that hypothesis-only baselines for NLI tasks can be remarkably robust. This question asks you to explore briefly how this baseline effects the 'edge_disjoint' and 'word_disjoint' versions of our task.

For this problem, submit two functions:

1. A `vector_combo_func` function called `hypothesis_only` that simply throws away the premise, using the unmodified hypothesis (second) vector as its representation of the example.

1. A function called `run_hypothesis_only_evaluation` that does the following:
    1. Loops over the two conditions 'word_disjoint' and 'edge_disjoint' and the two `vector_combo_func` values `vec_concatenate` and `hypothesis_only`, calling `nli.wordentail_experiment` to train on the conditions 'train' portion and assess on its 'dev' portion, with `glove_vec` as the `vector_func`. So that the results are consistent, use an `sklearn.linear_model.LogisticRegression` with default parameters as the model.
    1. Returns a `dict` mapping `(condition_name, function_name)` pairs to the 'macro-F1' score for that pair, as returned by the call to `nli.wordentail_experiment`. (Tip: you can get the `str` name of your function `hypothesis_only` with `hypothesis_only.__name__`.)

The test functions `test_hypothesis_only` and `test_run_hypothesis_only_evaluation` will help ensure that your functions have the desired logic.

```python
##### YOUR CODE HERE
from sklearn.linear_model import LogisticRegression

    ##### YOUR CODE HERE

def hypothesis_only(premise, hypothesis):
    #wordentail_data
    return hypothesis

#dict_keys(['edge_disjoint', 'vocab', 'word_disjoint'])

def run_hypothesis_only_evaluation():
    ##### YOUR CODE HERE
    condition_func_result = {}
    dataset_keys = ['word_disjoint', 'edge_disjoint']
    model_factory = lambda: LogisticRegression(fit_intercept=True, solver='liblinear')
    function_handles = [vec_concatenate, hypothesis_only]
    for dataset_key in dataset_keys:
        for vector_func in function_handles:
            splits = wordentail_data[dataset_key]

            result = nli.wordentail_experiment(
                    train_data=splits['train'],
                    assess_data=splits['dev'],
                    model=LogisticRegression(),
                    vector_func=glove_vec,
                    vector_combo_func=vec_concatenate)
            #print(result.keys())

            condition_func_result[(dataset_key,vector_func.__name__)] = result['macro-F1']

    return condition_func_result

```

```python
def test_hypothesis_only(hypothesis_only):
    v = hypothesis_only(1, 2)
    assert v == 2
```

```python
test_hypothesis_only(hypothesis_only)
```

```python
def test_run_hypothesis_only_evaluation(run_hypothesis_only_evaluation):
    results = run_hypothesis_only_evaluation()
    assert ('word_disjoint', 'vec_concatenate') in results, \
        "The return value of `run_hypothesis_only_evaluation` does not have the intended kind of keys"
    assert isinstance(results[('word_disjoint', 'vec_concatenate')], float), \
        "The values of the `run_hypothesis_only_evaluation` result should be floats"
```

```python
test_run_hypothesis_only_evaluation(run_hypothesis_only_evaluation)
```

                  precision    recall  f1-score   support

               0      0.901     0.980     0.939      1910
               1      0.466     0.142     0.218       239

        accuracy                          0.886      2149
       macro avg      0.684     0.561     0.578      2149
    weighted avg      0.853     0.886     0.859      2149

                  precision    recall  f1-score   support

               0      0.900     0.980     0.939      1910
               1      0.457     0.134     0.207       239

        accuracy                          0.886      2149
       macro avg      0.679     0.557     0.573      2149
    weighted avg      0.851     0.886     0.857      2149

                  precision    recall  f1-score   support

               0      0.875     0.969     0.920      7376
               1      0.574     0.229     0.328      1321

        accuracy                          0.857      8697
       macro avg      0.725     0.599     0.624      8697
    weighted avg      0.830     0.857     0.830      8697

                  precision    recall  f1-score   support

               0      0.875     0.970     0.920      7376
               1      0.573     0.224     0.322      1321

        accuracy                          0.857      8697
       macro avg      0.724     0.597     0.621      8697
    weighted avg      0.829     0.857     0.829      8697

### Alternatives to concatenation [2 points]

We've so far just used vector concatenation to represent the premise and hypothesis words. This question asks you to explore two simple alternative:

1. Write a function `vec_diff` that, for a given pair of vector inputs `u` and `v`, returns the element-wise difference between `u` and `v`.

1. Write a function `vec_max` that, for a given pair of vector inputs `u` and `v`, returns the element-wise max values between `u` and `v`.

You needn't include your uses of `nli.wordentail_experiment` with these functions, but we assume you'll be curious to see how they do!

```python
def vec_diff(u, v):
    ##### YOUR CODE HERE
    return u - v

def vec_max(u, v):
    ##### YOUR CODE HERE
    pass
    return np.maximum(u, v)

```

```python
def test_vec_diff(vec_diff):
    u = np.array([10.2, 8.1])
    v = np.array([1.2, -7.1])
    result = vec_diff(u, v)
    expected = np.array([9.0, 15.2])
    assert np.array_equal(result, expected), \
        "Expected {}; got {}".format(expected, result)
```

```python
test_vec_diff(vec_diff)
```

```python
def test_vec_max(vec_max):
    u = np.array([1.2,  8.1])
    v = np.array([10.2, -7.1])
    result = vec_max(u, v)
    expected = np.array([10.2, 8.1])
    assert np.array_equal(result, expected), \
        "Expected {}; got {}".format(expected, result)
```

```python
test_vec_max(vec_max)
```

### A deeper network [2 points]

It is very easy to subclass `TorchShallowNeuralClassifier` if all you want to do is change the network graph: all you have to do is write a new `define_graph`. If your graph has new arguments that the user might want to set, then you should also redefine `__init__` so that these values are accepted and set as attributes.

For this question, please subclass `TorchShallowNeuralClassifier` so that it defines the following graph:

$$\begin{align}
h_{1} &= xW_{1} + b_{1} \\
r_{1} &= \textbf{Bernoulli}(1 - \textbf{dropout\_prob}, n) \\
d_{1} &= r_1 * h_{1} \\
h_{2} &= f(d_{1}) \\
h_{3} &= h_{2}W_{2} + b_{2}
\end{align}$$

Here, $r_{1}$ and $d_{1}$ define a dropout layer: $r_{1}$ is a random binary vector of dimension $n$, where the probability of a value being $1$ is given by $1 - \textbf{dropout_prob}$. $r_{1}$ is multiplied element-wise by our first hidden representation, thereby zeroing out some of the values. The result is fed to the user's activation function $f$, and the result of that is fed through another linear layer to produce $h_{3}$. (Inside `TorchShallowNeuralClassifier`, $h_{3}$ is the basis for a softmax classifier, so no activation function is applied to it.)

For your implementation, please use `nn.Sequential`, `nn.Linear`, and `nn.Dropout` to define the required layers.

For comparison, using this notation, `TorchShallowNeuralClassifier` defines the following graph:

$$\begin{align}
h_{1} &= xW_{1} + b_{1} \\
h_{2} &= f(h_{1}) \\
h_{3} &= h_{2}W_{2} + b_{2}
\end{align}$$

The following code starts this sub-class for you, so that you can concentrate on `define_graph`. Be sure to make use of `self.dropout_prob`

For this problem, submit just your completed  `TorchDeepNeuralClassifier`. You needn't evaluate it, though we assume you will be keen to do that!

You can use `test_TorchDeepNeuralClassifier` to ensure that your network has the intended structure.

```python
import torch.nn as nn

class TorchDeepNeuralClassifier(TorchShallowNeuralClassifier):
    def __init__(self, dropout_prob=0.7, **kwargs):
        self.dropout_prob = dropout_prob
        super().__init__(**kwargs)

    def define_graph(self):
        """Complete this method!

        Returns
        -------
        an `nn.Module` instance, which can be a free-standing class you
        write yourself, as in `torch_rnn_classifier`, or the outpiut of
        `nn.Sequential`, as in `torch_shallow_neural_classifier`.

        """
        ##### YOUR CODE HERE
        return nn.Sequential(
            nn.Linear(self.input_dim, self.hidden_dim),
            nn.Dropout(p=self.dropout_prob),
            self.hidden_activation,
            nn.Linear(self.hidden_dim, self.n_classes_))

##### YOUR CODE HERE

net = TorchDeepNeuralClassifier(hidden_dim=50, max_iter=100)
word_disjoint_experiment = nli.wordentail_experiment(
    train_data=wordentail_data['word_disjoint']['train'],
    assess_data=wordentail_data['word_disjoint']['dev'],
    model=net,
    vector_func=glove_vec,
    vector_combo_func=vec_concatenate)

```

    Finished epoch 100 of 100; error is 2.4060109853744507

                  precision    recall  f1-score   support

               0      0.900     0.995     0.945      1910
               1      0.737     0.117     0.202       239

        accuracy                          0.897      2149
       macro avg      0.818     0.556     0.574      2149
    weighted avg      0.882     0.897     0.862      2149

```python
def test_TorchDeepNeuralClassifier(TorchDeepNeuralClassifier):
    dropout_prob = 0.55
    assert hasattr(TorchDeepNeuralClassifier(), "dropout_prob"), \
        "TorchDeepNeuralClassifier must have an attribute `dropout_prob`."
    try:
        inst = TorchDeepNeuralClassifier(dropout_prob=dropout_prob)
    except TypeError:
        raise TypeError("TorchDeepNeuralClassifier must allow the user "
                        "to set `dropout_prob` on initialization")
    inst.input_dim = 10
    inst.n_classes_ = 5
    graph = inst.define_graph()
    assert len(graph) == 4, \
        "The graph should have 4 layers; yours has {}".format(len(graph))
    expected = {
        0: 'Linear',
        1: 'Dropout',
        2: 'Tanh',
        3: 'Linear'}
    for i, label in expected.items():
        name = graph[i].__class__.__name__
        assert label in name, \
            "The {} layer of the graph should be a {} layer; yours is {}".format(i, label, name)
    assert graph[1].p == dropout_prob, \
        "The user's value for `dropout_prob` should be the value of `p` for the Dropout layer."
```

```python
test_TorchDeepNeuralClassifier(TorchDeepNeuralClassifier)
```

### Your original system [3 points]

This is a simple dataset, but our focus on the 'word_disjoint' condition ensures that it's a challenging one, and there are lots of modeling strategies one might adopt.

You are free to do whatever you like. We require only that your system differ in some way from those defined in the preceding questions. They don't have to be completely different, though. For example, you might want to stick with the model but represent examples differently, or the reverse.

Keep in mind that, for the bake-off evaluation, the 'edge_disjoint' portions of the data are off limits. You can, though, train on the combination of the 'word_disjoint' 'train' and 'dev' portions. You are free to use different pretrained word vectors and the like. Please do not introduce additional entailment datasets into your training data, though.

Please embed your code in this notebook so that we can rerun it.

In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies.

```python
# Enter your system description in this cell.
"""Baseline for quick lookup
              precision    recall  f1-score   support

           0      0.924     0.938     0.931      1910
           1      0.435     0.381     0.406       239

    accuracy                          0.876      2149
   macro avg      0.680     0.659     0.669      2149
weighted avg      0.869     0.876     0.873      2149

#https://github.com/facebookresearch/fastText/issues/458
#https://radimrehurek.com/gensim/models/fasttext.html
#https://radimrehurek.com/gensim/models/fasttext.html#gensim.models.fasttext.load_facebook_model
#Enriching Word Vectors with Subword Information https://arxiv.org/abs/1607.04606

The baseline does well to begin with.
But it has a flaw that  OOV (out of vocabulary)  words are assigned random vectors.
I thought Facebook fasttext+genism would alleviate the issue because they are trained on word morphology as well.
By using embeddings from above the performance still slightly fell short of baselines even after a
grid search for best hidden layer size.

Next thing I tried was Glove-300 with retrofitting. Retrofitting does 0 iterations which means we might already
have high quality vectors. However inorder to solve the problem of out of vocabulary words
I used the approach in the paper `Enriching Word Vectors with Subword Information`.
i.e. take n-grams of length 3-6 from words use their embeddings for enrichment.
SO for a OOV word I do a mean of the ngrams of the above specified length.
If no n-gram is found I just return a random vector

Since the system didn't beat the baseline, I tried to get better embeddings. I used glove 840B common crawl
with 300 dimensions. This approach again uses similar logic as above to handle OOV words.
Further when running a hyperparameter search I realized that the system doesn't improve much beyond 100 iterations.

I also tried multiple network depths but finally only selected to add an additional linear layer and activation i.e.
DeepClassifier class defined below. So this makes the final pipeline as a Deepclassifier using concatenate function.
Using ngram decomposition to get representation for missing words.

My peak score was: 0.63

 """

# Please do not remove this comment.

if 'IS_GRADESCOPE_ENV' not in os.environ:
    from gensim.test.utils import datapath
    from gensim.models.fasttext import load_facebook_model
    from matplotlib.pyplot import plot
    from collections import defaultdict
    from nltk.corpus import wordnet as wn
    import numpy as np
    import os
    import pandas as pd
    import retrofitting
    from retrofitting import Retrofitter
    import utils

    def load_fast_text_vectors(fname):
        fin = io.open(fname, 'r', encoding='utf-8', newline='\n', errors='ignore')
        n, d = map(int, fin.readline().split())
        data = {}
        for line in fin:
            tokens = line.rstrip().split(' ')
            data[tokens[0]] = map(float, tokens[1:])
        return data

    #def load_vie_gen(bin_path="/Users/simranjit/Neuron/repos/cs22u_fork/data/nlidata/crawl-300d-2M-subword/crawl-300d-2M-subword.bin"):
    def load_vie_gen(bin_path=os.path.join(DATA_HOME, "crawl-300d-2M-subword/crawl-300d-2M-subword.bin")):
        #if not bin_path:
        #    bin_path =
        cap_path = datapath(bin_path)
        fb_model = load_facebook_model(cap_path)
        return fb_model

    #ftt_model = load_vie_gen()
    def ftt_vec(word):
        return ftt_model[word]

    def load_glove_any(my_dims, tokens):
        GLOVE_HOME_HERE = os.path.join(DATA_HOME, 'glove.{}B'.format(tokens))
        glove_src_mine = os.path.join(GLOVE_HOME_HERE, 'glove.{}B.{}d.txt'.format(tokens, my_dims))
        GLOVE_mine = utils.glove2dict(glove_src_mine)
        return GLOVE_mine

    # Creates a dict mapping strings (words) to GloVe vectors:
    my_glove_dim = 300
    token_count = 6
    #GLOVE_mine = load_glove_any(my_glove_dim, token_count)
    token_count = 840
    GLOVE_mine = load_glove_any(my_glove_dim, 840)
   ###############
    def extract_ngrams(word):
        #print(word)
        if (len(word)<=2 ):
            return []
        n = len(word)
        ngrams = []
        for i in range (3, 7):
            for j in range(0, n-i):
                gram = word[j:j+i]
                ngrams.append(gram)
        return ngrams

    def get_unknown_grams(word, embedding_dict):
        ngrams = extract_ngrams(word)
        if ngrams is None:
            print(word)
        embeddings = []
        #print(ngrams)
        for gram in ngrams:
            if gram in embedding_dict:
                embeddings.append(embedding_dict[gram])
        #print(embeddings)
        return embeddings

    def glove_vec_mine(w):
        """Return `w`'s GloVe representation if available, else return
        a random vector."""
        #if w not in GLOVE:
        #print(w)
        if w in GLOVE_mine:
            return GLOVE_mine[w]

        ngram_embeddings = get_unknown_grams(w, GLOVE_mine)

        if len(ngram_embeddings)==0:
            return randvec(w, n=my_glove_dim)
        fin = np.mean(ngram_embeddings, axis=0)
        #print(fin)
        return fin

    ####################
    def get_wordnet_edges():
        edges = defaultdict(set)
        for ss in wn.all_synsets():
            lem_names = {lem.name() for lem in ss.lemmas()}
            for lem in lem_names:
                edges[lem] |= lem_names
        return edges

    def convert_edges_to_indices(edges, Q):
        lookup = dict(zip(Q.index, range(Q.shape[0])))
        index_edges = defaultdict(set)
        for start, finish_nodes in edges.items():
            s = lookup.get(start)
            if s:
                f = {lookup[n] for n in finish_nodes if n in lookup}
                if f:
                    index_edges[s] = f
        return index_edges

    #wn_edges = get_wordnet_edges()

    """
    def run_hyper_param_search(vector_func, vector_combo_func):
        metric_key = "macro-F1"
        dim_list = []
        score_list = []
        #base_score = word_disjoint_experiment[metric_key]
        base_score = 0.67
        for dims in range(1,26):
            final_dims = 50+dims*10
            net = TorchShallowNeuralClassifier(hidden_dim=final_dims, max_iter=100)
            word_disjoint_experiment = nli.wordentail_experiment(
                train_data=wordentail_data['word_disjoint']['train'],
                assess_data=wordentail_data['word_disjoint']['dev'],
                model=net,
                vector_func=vector_func,
                vector_combo_func=vector_combo_func)
            score = word_disjoint_experiment[metric_key]
            if score > base_score:
                ("following pair beat it")
            print(final_dims, score)
            dim_list.append(final_dims)
            score_list.append(score)
        plot(dim_list, score_list)
        return dim_list, score_list
        """

    def retrofit_for_dict(embedding_dict):
        embedding_df = pd.DataFrame(list(embedding_dict.values()))
        # print(embedding_df)
        wn_index_edges = convert_edges_to_indices(wn_edges, embedding_df)
        fitter = Retrofitter(verbose=True)
        wn_retro = fitter.fit(embedding_df, wn_index_edges)
        return wn_retro.to_numpy()

    #retrofitted = retrofit_for_dict(GLOVE_mine)

```

```python

```

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:

    import torch
    import torch.nn as nn

    class DeepClassifier(TorchShallowNeuralClassifier):
        def __init__(self, **kwargs):
            super(DeepClassifier, self).__init__(**kwargs)

        def define_graph(self):
            return nn.Sequential(
                nn.Linear(self.input_dim, self.hidden_dim),
                self.hidden_activation,

                #nn.Linear(self.hidden_dim, self.hidden_dim),
                #self.hidden_activation,

                nn.Linear(self.hidden_dim, self.hidden_dim),
                self.hidden_activation,

                nn.Linear(self.hidden_dim, self.n_classes_))

    def run_hyper_param_search(vector_func, vector_combo_func):
            metric_key = "macro-F1"
            dim_list = []
            score_list = []
            iters_list = []
            #base_score = word_disjoint_experiment[metric_key]
            base_score = 0.67

            start_iters = 120

            iters = start_iters
            dims = 80

            if True:
            #for iters in range(start_iters, 201, 10):
            #for dims in range(1,26):
            #for dims in range(50, 150, 10):
                #final_dims = 50+dims*10
                #final_dims = 150
                #net = TorchShallowNeuralClassifier(hidden_dim=final_dims, max_iter=iters)
                final_dims = dims
                net = DeepClassifier(hidden_dim=final_dims, max_iter=iters)
                word_disjoint_experiment = nli.wordentail_experiment(
                    train_data=wordentail_data['word_disjoint']['train'],
                    assess_data=wordentail_data['word_disjoint']['dev'],
                    model=net,
                    vector_func=vector_func,
                    vector_combo_func=vector_combo_func)
                score = word_disjoint_experiment[metric_key]
                if score > base_score:
                    ("following pair beat it")
                print(final_dims, score)
                dim_list.append(final_dims)
                score_list.append(score)
                iters_list.append(iters)

            list_to_plot = dim_list
            if (dim_list[0] == dim_list[-1]):
                list_to_plot = iters_list

            plot(list_to_plot, score_list)
            return list_to_plot, score_list

    run_hyper_param_search(glove_vec_mine, vec_concatenate)
    """
    net = TorchDeepNeuralClassifier(hidden_dim=300, max_iter=100)
    word_disjoint_experiment = nli.wordentail_experiment(
        train_data=wordentail_data['word_disjoint']['train'],
        assess_data=wordentail_data['word_disjoint']['dev'],
        model=net,
        vector_func=glove_vec_mine,
        vector_combo_func=vec_concatenate)

    """
```

```python

```

## Bake-off [1 point]

The goal of the bake-off is to achieve the highest macro-average F1 score on __word_disjoint__, on a test set that we will make available at the start of the bake-off. The announcement will go out on the discussion forum. To enter, you'll be asked to run `nli.bake_off_evaluation` on the output of your chosen `nli.wordentail_experiment` run.

The cells below this one constitute your bake-off entry.

The rules described in the [Your original system](#Your-original-system-[3-points]) homework question are also in effect for the bake-off.

Systems that enter will receive the additional homework point, and systems that achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.

Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.

The announcement will include the details on where to submit your entry.

```python
# Enter your bake-off assessment code into this cell.
# Please do not remove this comment.
##### YOUR CODE HERE
#print(net)
final_dims = 80
iters = 120
entailment_model = DeepClassifier(hidden_dim=final_dims, max_iter=iters)
bakeoff_model = nli.wordentail_experiment(
    train_data=wordentail_data['word_disjoint']['train'],
    assess_data=wordentail_data['word_disjoint']['dev'],
    model=entailment_model,
    vector_func=glove_vec_mine,
    vector_combo_func=vec_concatenate)

test_data_filename = os.path.join(
    NLIDATA_HOME,
    "bakeoff-wordentail-data",
    "nli_wordentail_bakeoff_data-test.json")

nli.bake_off_evaluation(
    bakeoff_model,
    test_data_filename)

```

    Finished epoch 120 of 120; error is 0.00010514869154576445

                  precision    recall  f1-score   support

               0      0.920     0.924     0.922      1910
               1      0.372     0.360     0.366       239

        accuracy                          0.861      2149
       macro avg      0.646     0.642     0.644      2149
    weighted avg      0.859     0.861     0.860      2149

                  precision    recall  f1-score   support

               0      0.870     0.896     0.883      1767
               1      0.533     0.471     0.500       446

        accuracy                          0.810      2213
       macro avg      0.702     0.683     0.691      2213
    weighted avg      0.802     0.810     0.806      2213

```python
# On an otherwise blank line in this cell, please enter
# your macro-avg f1 value as reported by the code above.
# Please enter only a number between 0 and 1 inclusive.
# Please do not remove this comment.

##### YOUR CODE HERE
0.691

```


# snapReactors

Copyright (c) Dan Kotlyar and CoRE group

# Property Container

* This container stores all relevant info required for a specific
property and provides methods to evaluate propertys based on specfic
dependencies such as temperature, pressure, etc.
* The container also organizes propertys based on their type, providing subclasses such
as Constant, Table, and Correlation.

### Code:

```python
    >>> p2 = Table('h', np.array([1, 2, 3, 4]), 'W/K*m^2',
    >>>     np.array([100, 200, 300, 400]), 'K',
    >>>     unc = np.array([.01, .01, .01, .01]))
    >>>
    >>> p2.evaluate(300) #300K
```

### Notes:

* This snippet shows the use of the tabulated property container ``Property.Table``.
* Property values can be given in tabulated fashion with up to two dependencies.
* The ``Property.evaluate(dependecys)`` method can be used to evaluate the method
for different dependency values.

```python
import numpy as np
import sympy as sp
from snapReactors.containers.property import Property, Constant, Table, Correlation
```

#### Defining a constant property
1. Give name of property, must be in ``ALLOWED_PROPERTIES``
3. Give value of property
4. Give units of property
5. value uncertainty, ref, and description are left as optional parameters

```python
p1 = Constant(id='cv',  value=1, unit= "J/kg/K", unc=None, ref=None, description=None)
print(p1)
```

    {'id': 'cv', 'dtype': <DTYPE.NUMBER: 1>, 'vtype': <VTYPE.CONSTANT: 1>, 'value': array([1]), 'valueUnit': 'J/kg/K', 'unc': None, 'dependents': None, 'dependentsUnit': None, 'description': None, 'ref': None}

#### Defining a tabulated property
1. Give name of property, must be in ``ALLOWED_PROPERTIES``
3. Give values of property
4. Give units of property
5. Give values of 1st dependency
6. Give units of 1st dependency
7. 2nd dependency values, 2nd dependency units, value uncertainty, reference, and description are left as optional parameters

```python
p2 = Table('h',  np.array([1, 2, 3, 4]), 'W/K/m^2', np.array([100, 200, 300, 400]), 'K',  unc = np.array([.01, .01, .01, .01]), dependency2=None, dependencyUnit2=None, ref=None, description=None)
print(p2)
```

    {'id': 'h', 'dtype': <DTYPE.NDARRAY: 2>, 'vtype': <VTYPE.TABLE: 2>, 'value': array([1, 2, 3, 4]), 'valueUnit': 'W/K/m^2', 'unc': array([0.01, 0.01, 0.01, 0.01]), 'dependents': array([100, 200, 300, 400]), 'dependentsUnit': 'K', 'description': None, 'ref': None, 'dependency1': array([100, 200, 300, 400]), 'dependency2': None, 'dependencyUnit1': 'K', 'dependencyUnit2': None}

#### Defining a correlation based property
1. Give name of property, must be in ``ALLOWED_PROPERTIES``
3. Give string expression for correlation, follows python syntax
4. Give string of symbols for expression, symbols delineated by commas
4. Give units of property
5. Give dependency range, bounds for the correlation
6. Give units of dependency
7. value uncertainty, reference, and description are left as optional parameters

```python
corr1 = "T**2 + P + 1/2"
syms1 = "T, P"

p3 = Correlation('h', corr1, syms1, 'W/K/m^2', np.array([300, 600]), 'K', np.array([10, 20]), 'Pa', unc=None, ref=None, description=None)
print(p3)
```

    {'id': 'h', 'dtype': <DTYPE.NDARRAY: 2>, 'vtype': <VTYPE.CORRELATION: 3>, 'value': array([ 90010.5       ,  93721.6578509 ,  97507.78446481, 101368.87984173,
           105304.94398167, 109315.97688463, 113401.9785506 , 117562.94897959,
           121798.8881716 , 126109.79612661, 130495.67284465, 134956.5183257 ,
           139492.33256976, 144103.11557684, 148788.86734694, 153549.58788005,
           158385.27717618, 163295.93523532, 168281.56205748, 173342.15764265,
           178477.72199084, 183688.25510204, 188973.75697626, 194334.22761349,
           199769.66701374, 205280.07517701, 210865.45210329, 216525.79779259,
           222261.1122449 , 228071.39546022, 233956.64743857, 239916.86817993,
           245952.0576843 , 252062.21595169, 258247.34298209, 264507.43877551,
           270842.50333195, 277252.5366514 , 283737.53873386, 290297.50957934,
           296932.44918784, 303642.35755935, 310427.23469388, 317287.08059142,
           324221.89525198, 331231.67867555, 338316.43086214, 345476.15181175,
           352710.84152436, 360020.5       ]), 'valueUnit': 'W/K/m^2', 'unc': None, 'dependents': array([[300.        , 306.12244898, 312.24489796, 318.36734694,
            324.48979592, 330.6122449 , 336.73469388, 342.85714286,
            348.97959184, 355.10204082, 361.2244898 , 367.34693878,
            373.46938776, 379.59183673, 385.71428571, 391.83673469,
            397.95918367, 404.08163265, 410.20408163, 416.32653061,
            422.44897959, 428.57142857, 434.69387755, 440.81632653,
            446.93877551, 453.06122449, 459.18367347, 465.30612245,
            471.42857143, 477.55102041, 483.67346939, 489.79591837,
            495.91836735, 502.04081633, 508.16326531, 514.28571429,
            520.40816327, 526.53061224, 532.65306122, 538.7755102 ,
            544.89795918, 551.02040816, 557.14285714, 563.26530612,
            569.3877551 , 575.51020408, 581.63265306, 587.75510204,
            593.87755102, 600.        ],
           [ 10.        ,  10.20408163,  10.40816327,  10.6122449 ,
             10.81632653,  11.02040816,  11.2244898 ,  11.42857143,
             11.63265306,  11.83673469,  12.04081633,  12.24489796,
             12.44897959,  12.65306122,  12.85714286,  13.06122449,
             13.26530612,  13.46938776,  13.67346939,  13.87755102,
             14.08163265,  14.28571429,  14.48979592,  14.69387755,
             14.89795918,  15.10204082,  15.30612245,  15.51020408,
             15.71428571,  15.91836735,  16.12244898,  16.32653061,
             16.53061224,  16.73469388,  16.93877551,  17.14285714,
             17.34693878,  17.55102041,  17.75510204,  17.95918367,
             18.16326531,  18.36734694,  18.57142857,  18.7755102 ,
             18.97959184,  19.18367347,  19.3877551 ,  19.59183673,
             19.79591837,  20.        ]]), 'dependentsUnit': 'K, Pa', 'description': None, 'ref': None, 'expr': 'T**2 + P + 1/2', 'syms': 'T, P', 'dependencyRange1': array([300, 600]), 'dependencyUnit1': 'K', 'dependencyRange2': array([10, 20]), 'dependencyUnit2': 'Pa'}

### Evaluating Propertys at specified dependencys
1. The ``Property.evaluate(dependencys)`` method is used to evaluate propertys at specfied dependencys
2. The evaluate method distinguish between propertys based on their type
3. The value returned by the method will be different evaluated differently for each type.

#### Evaluating Constants at specified dependencys
1. Constant will remain constant regardless of dependencies

```python
v1p1 = p1.evaluate(300, 1500)
v2p1 = p1.evaluate(400, 1500)

print(v1p1, v2p1)

```

    [1] [1]

#### Evaluating Correlation propertys at specified dependencys
1. Correlations only support 2 dependency and will output the evualted value at the given dependency/s

```python
v1p3 = p3.evaluate(dependency1=300, dependency2=15)
v2p3 = p3.evaluate(dependency1=400, dependency2=15)

print(v1p3, v2p3)
```

    90015.5000000000 160015.500000000

#### Evaluating Tabulated propertys at specified dependencys
1. Tabulated propertys only support 2 dependencys for now and will output the evaluated value at the given dependencys based off a linear interpolation method

```python
v1p2 = p2.evaluate(dependency1=257)
v2p2 = p2.evaluate(dependency1=143)

print(v1p2, v2p2)
```

    2.57 1.43


## Contents

* amortized variational inference (VAEs)
* improving amortized variational inference using KL annleaing
* improving amortized variational inference using constraint optimization
* avoiding latent space distribution matching using GANs

```python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import tensorflow as tf
import tensorflow_datasets as tfds

import numpy as np

# Plotting library.
from matplotlib import pyplot as plt
import seaborn as sns

# Tensorflow probability utilities
import tensorflow_probability as tfp

tfd = tfp.distributions

import warnings

```

```python
sns.set(rc={"lines.linewidth": 2.8}, font_scale=2)
sns.set_style("whitegrid")
```

```python
# # Don't forget to select GPU runtime environment in Runtime -> Change runtime type
# device_name = tf.test.gpu_device_name()
# if device_name != '/device:GPU:0':
#   raise SystemError('GPU device not found')
# print('Found GPU at: {}'.format(device_name))
```

```python
def gallery(array, ncols=10, rescale=False):
    """Data visualization code."""
    if rescale:
      array = (array + 1.) / 2
    nindex, height, width, intensity = array.shape
    nrows = nindex//ncols
    assert nindex == nrows*ncols
    # want result.shape = (height*nrows, width*ncols, intensity)
    result = (array.reshape(nrows, ncols, height, width, intensity)
              .swapaxes(1,2)
              .reshape(height*nrows, width*ncols, intensity))
    return result
```

```python
def show_digits(axis, digits, title=''):
  axis.axis('off')
  ncols = int(np.sqrt(digits.shape[0]))
  axis.imshow(gallery(digits, ncols=ncols).squeeze(axis=2),
                 cmap='gray')
  axis.set_title(title, fontsize=15)
```

```python
def show_latent_interpolations(generator, prior, session):
  a = np.linspace(0.0, 1.0, BATCH_SIZE)
  a = np.expand_dims(a, axis=1)

  first_latents = prior.sample()[0]
  second_latents = prior.sample()[0]

  # To ensure that the interpolation is still likely under the Gaussian prior,
  # we use Gaussian interpolation - rather than linear interpolation.
  interpolations =  np.sqrt(a) * first_latents + np.sqrt(1 - a) * second_latents

  ncols = int(np.sqrt(BATCH_SIZE))
  samples_from_interpolations = generator(interpolations)
  samples_from_interpolations_np = sess.run(samples_from_interpolations)
  plt.gray()
  axis = plt.gca()
  show_digits(
     axis, samples_from_interpolations_np, title='Latent space interpolations')
```

```python
BATCH_SIZE = 64
NUM_LATENTS = 10
TRAINING_STEPS = 10000
NUM_SVI_UPDATES = 100
```

```python
DECODER_VARIABLE_SCOPE = "decoder"
```

```python
tf.reset_default_graph()
```

## Get the Data

We will use the MNIST dataset. Luckly, TensorFlow comes with a simple way to load it.

```python
warnings.filterwarnings('ignore')
mnist = tf.contrib.learn.datasets.load_dataset("mnist")
print(mnist.train.images.shape)
print(type(mnist.train.images))
```

    WARNING: Logging before flag parsing goes to stderr.
    W0715 15:38:35.874579 4672783808 lazy_loader.py:50]
    The TensorFlow contrib module will not be included in TensorFlow 2.0.
    For more information, please see:
      * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
      * https://github.com/tensorflow/addons
      * https://github.com/tensorflow/io (for I/O related ops)
    If you depend on functionality not listed there, please file an issue.

    W0715 15:38:35.875852 4672783808 deprecation.py:323] From <ipython-input-11-da699ef53d46>:2: load_dataset (from tensorflow.contrib.learn.python.learn.datasets) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use tf.data.
    W0715 15:38:35.876689 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/__init__.py:80: load_mnist (from tensorflow.contrib.learn.python.learn.datasets.mnist) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use alternatives such as official/mnist/dataset.py from tensorflow/models.
    W0715 15:38:35.879344 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/mnist.py:300: read_data_sets (from tensorflow.contrib.learn.python.learn.datasets.mnist) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use alternatives such as official/mnist/dataset.py from tensorflow/models.
    W0715 15:38:35.880613 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/mnist.py:260: maybe_download (from tensorflow.contrib.learn.python.learn.datasets.base) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please write your own downloading logic.
    W0715 15:38:35.882920 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/base.py:252: _internal_retry.<locals>.wrap.<locals>.wrapped_fn (from tensorflow.contrib.learn.python.learn.datasets.base) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use urllib or similar directly.
    W0715 15:39:02.161442 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/mnist.py:262: extract_images (from tensorflow.contrib.learn.python.learn.datasets.mnist) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use tf.data to implement this functionality.

    Successfully downloaded train-images-idx3-ubyte.gz 9912422 bytes.
    Extracting MNIST-data/train-images-idx3-ubyte.gz

    W0715 15:39:02.602097 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/mnist.py:267: extract_labels (from tensorflow.contrib.learn.python.learn.datasets.mnist) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use tf.data to implement this functionality.

    Successfully downloaded train-labels-idx1-ubyte.gz 28881 bytes.
    Extracting MNIST-data/train-labels-idx1-ubyte.gz

    W0715 15:39:05.039743 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/contrib/learn/python/learn/datasets/mnist.py:290: DataSet.__init__ (from tensorflow.contrib.learn.python.learn.datasets.mnist) is deprecated and will be removed in a future version.
    Instructions for updating:
    Please use alternatives such as official/mnist/dataset.py from tensorflow/models.

    Successfully downloaded t10k-images-idx3-ubyte.gz 1648877 bytes.
    Extracting MNIST-data/t10k-images-idx3-ubyte.gz
    Successfully downloaded t10k-labels-idx1-ubyte.gz 4542 bytes.
    Extracting MNIST-data/t10k-labels-idx1-ubyte.gz
    (55000, 784)
    <class 'numpy.ndarray'>

```python
print(mnist.train.images.shape)
print(type(mnist.train.images))
```

    (55000, 784)
    <class 'numpy.ndarray'>

### Transform the Data from Numpy Arrays to in Graph Tensors.

This allows us to use TensorFlow datasets, which ensure that a new batch from the data is being fed at each session.run. This means that we do not need to use feed_dicts to feed data to each session.

```python
def make_tf_data_batch(np_data, shuffle=True):
  # Reshape the data to image size.
  images = np_data.reshape((-1, 28, 28, 1))

  # Create the TF dataset.
  dataset = tf.data.Dataset.from_tensor_slices(images)

  # Shuffle and repeat the dataset for training.
  # This is required because we want to do multiple passes through the entire
  # dataset when training.
  if shuffle:
    dataset = dataset.shuffle(100000).repeat()

  # Batch the data and return the data batch.
  batch_data = dataset.batch(BATCH_SIZE)
  one_shot_iterator = tf.compat.v1.data.make_one_shot_iterator(batch_data)
  data_batch = one_shot_iterator.get_next()

  return data_batch
```

```python
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

    (?, 28, 28, 1)

# Amortized variational inference (VAEs)

Instead of learning one set of posterior variables per data point, we can use function approximation to learn the distributional variables. Specifically, the posterior parameters for $x_i$ will be the output of a *learned* function $f_\theta(x_i)$, where $\theta$ are parameters shared across all data points. Can you think of why this is useful?

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)}  - \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

For more information, see:
  * https://arxiv.org/abs/1312.6114

```python
def standard_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```python
def big_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```python
ENCODER_VARIABLE_SCOPE = 'encoder'
```

```python
def encoder(x):
  with tf.variable_scope(ENCODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):

    h = tf.layers.Conv2D(
          filters=8,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(x)
    h = tf.layers.Conv2D(
          filters=16,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2D(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)

    out_shape = 1
    for s in h.shape.as_list()[1:]:
      out_shape*= s

    h = tf.reshape(h, shape=[BATCH_SIZE, out_shape])
    mean = tf.layers.dense(h, NUM_LATENTS, activation=None)
    scale = tf.layers.dense(h, NUM_LATENTS, activation=None)
    return multi_normal(loc=mean, log_scale=scale)
```

```python
def multi_normal(loc, log_scale):
  # We model the latent variables as independent
  return tfd.Independent(
      distribution=tfd.Normal(loc=loc, scale=tf.exp(log_scale)),
      reinterpreted_batch_ndims=1)
```

```python
def make_prior():
  # Zero mean, unit variance prior.
  prior_mean = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)
  prior_log_scale = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)

  return multi_normal(prior_mean, prior_log_scale)
```

```python
def bound_terms(data_batch, variational_posterior, decoder_fn):
  latent_sample = variational_posterior.sample()
  all_log_probs = decoder_fn(latent_sample).log_prob(data_batch)

  # Since we used broadcasting, log_prob of a Bernoulli will return a log_prob for
  # each dimension.
  all_log_probs.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

  # Reduce sum over the data dimensions
  likelihood_term = tf.reduce_sum(all_log_probs, axis=[1, 2, 3])

  # Reduce mean over the batch dimensions
  likelihood_term = tf.reduce_mean(likelihood_term)

  kl_term = variational_posterior.kl_divergence(prior)
  kl_term.shape.assert_is_compatible_with([BATCH_SIZE])

  # Reduce over the batch dimension.
  kl_term = tf.reduce_mean(kl_term)

  return likelihood_term, kl_term
```

```python
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

    (?, 28, 28, 1)

```python
prior = make_prior()
```

```python
# Note: we no longer have to use a variable to store the data.
# We will perform one encoder update per decoder update, so it is OK for the
# data batch to be refreshed at each run.
variational_posterior = encoder(real_data)
decoder = standard_decoder
```

    W0715 15:40:14.153403 4672783808 deprecation.py:506] From /anaconda3/lib/python3.7/site-packages/tensorflow/python/ops/init_ops.py:1251: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
    Instructions for updating:
    Call initializer instance with the dtype argument instead of passing it to the constructor
    W0715 15:40:14.467746 4672783808 deprecation.py:323] From <ipython-input-19-997b61b0c8ff>:28: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.
    Instructions for updating:
    Use keras.layers.dense instead.

```python
# Maximize the data likelihodd and minimize the KL divergence between the prior
# and posterior. We use the exact same loss as in the SVI case.
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
train_elbo = likelihood_term - kl_term
loss =  - train_elbo
```

    W0715 15:40:16.297658 4672783808 deprecation.py:323] From <ipython-input-15-a043c76715d6>:17: Bernoulli.__init__ (from tensorflow.python.ops.distributions.bernoulli) is deprecated and will be removed after 2019-01-01.
    Instructions for updating:
    The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.distributions`.
    W0715 15:40:16.300471 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/python/ops/distributions/bernoulli.py:97: Distribution.__init__ (from tensorflow.python.ops.distributions.distribution) is deprecated and will be removed after 2019-01-01.
    Instructions for updating:
    The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.distributions`.
    W0715 15:40:16.319782 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/python/ops/nn_impl.py:180: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
    Instructions for updating:
    Use tf.where in 2.0, which has the same broadcast rule as np.where

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
update_op = optimizer.minimize(loss, var_list=autoencoder_vars)
```

## Training

```python
sess = tf.Session()

# Initialize all variables
# sess.run(tf.initialize_all_variables())
sess.run(tf.global_variables_initializer())
```

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

    Iteration 0. Loss 540.9539794921875. KL 0.008539748378098011
    Iteration 100. Loss 204.42613220214844. KL 6.141989707946777
    Iteration 200. Loss 182.2940673828125. KL 6.860536098480225
    Iteration 300. Loss 143.0342559814453. KL 13.470229148864746
    Iteration 400. Loss 128.14712524414062. KL 15.409025192260742
    Iteration 500. Loss 119.27674865722656. KL 16.24714469909668
    Iteration 600. Loss 123.24623107910156. KL 17.276906967163086
    Iteration 700. Loss 116.22936248779297. KL 17.678213119506836
    Iteration 800. Loss 121.54521179199219. KL 17.341318130493164
    Iteration 900. Loss 118.71582794189453. KL 16.90508270263672
    Iteration 1000. Loss 115.46253204345703. KL 17.668479919433594
    Iteration 1100. Loss 118.37068176269531. KL 18.00971031188965
    Iteration 1200. Loss 111.82579040527344. KL 17.66265869140625
    Iteration 1300. Loss 112.63470458984375. KL 17.121410369873047
    Iteration 1400. Loss 104.64066314697266. KL 17.39837646484375
    Iteration 1500. Loss 108.8840103149414. KL 17.932090759277344
    Iteration 1600. Loss 106.60586547851562. KL 17.796913146972656
    Iteration 1700. Loss 113.857421875. KL 17.66264533996582
    Iteration 1800. Loss 108.50210571289062. KL 19.06524658203125
    Iteration 1900. Loss 112.80574035644531. KL 19.243526458740234
    Iteration 2000. Loss 116.29713439941406. KL 18.42117691040039
    Iteration 2100. Loss 118.50651550292969. KL 19.16826057434082
    Iteration 2200. Loss 106.24510192871094. KL 19.17131233215332
    Iteration 2300. Loss 106.39988708496094. KL 18.242839813232422
    Iteration 2400. Loss 106.77734375. KL 17.94028091430664
    Iteration 2500. Loss 104.30854797363281. KL 19.087722778320312
    Iteration 2600. Loss 104.84104919433594. KL 18.804929733276367
    Iteration 2700. Loss 104.9285888671875. KL 18.637866973876953
    Iteration 2800. Loss 106.45011901855469. KL 18.45279312133789
    Iteration 2900. Loss 105.27384185791016. KL 17.863876342773438
    Iteration 3000. Loss 111.72129821777344. KL 18.84537124633789
    Iteration 3100. Loss 109.09244537353516. KL 19.24643898010254
    Iteration 3200. Loss 108.33090209960938. KL 19.440717697143555
    Iteration 3300. Loss 104.71376037597656. KL 19.263229370117188
    Iteration 3400. Loss 99.09135437011719. KL 18.66463851928711
    Iteration 3500. Loss 105.83938598632812. KL 18.857175827026367
    Iteration 3600. Loss 111.00079345703125. KL 19.117870330810547
    Iteration 3700. Loss 109.00234985351562. KL 18.935073852539062
    Iteration 3800. Loss 108.42906951904297. KL 18.956382751464844
    Iteration 3900. Loss 96.23290252685547. KL 18.47344207763672
    Iteration 4000. Loss 104.66486358642578. KL 18.764307022094727
    Iteration 4100. Loss 102.00324249267578. KL 18.908668518066406
    Iteration 4200. Loss 109.47421264648438. KL 18.88107681274414
    Iteration 4300. Loss 101.82896423339844. KL 18.277616500854492
    Iteration 4400. Loss 106.3016357421875. KL 18.90798568725586
    Iteration 4500. Loss 106.20803833007812. KL 19.361387252807617
    Iteration 4600. Loss 103.95289611816406. KL 19.06844711303711
    Iteration 4700. Loss 99.21043395996094. KL 18.777164459228516
    Iteration 4800. Loss 101.5719985961914. KL 18.647300720214844
    Iteration 4900. Loss 105.81336975097656. KL 18.95328140258789
    Iteration 5000. Loss 104.25371551513672. KL 18.61078643798828
    Iteration 5100. Loss 99.74555206298828. KL 19.20177459716797
    Iteration 5200. Loss 106.2641372680664. KL 19.449623107910156
    Iteration 5300. Loss 102.13298034667969. KL 20.386205673217773
    Iteration 5400. Loss 99.3668441772461. KL 19.57019805908203
    Iteration 5500. Loss 106.29104614257812. KL 19.370845794677734
    Iteration 5600. Loss 111.29519653320312. KL 20.31646728515625
    Iteration 5700. Loss 108.01383972167969. KL 20.06625747680664
    Iteration 5800. Loss 100.60852813720703. KL 18.658178329467773
    Iteration 5900. Loss 104.4432373046875. KL 18.553207397460938
    Iteration 6000. Loss 100.87751770019531. KL 18.69692611694336
    Iteration 6100. Loss 109.80620574951172. KL 19.759727478027344
    Iteration 6200. Loss 100.78556823730469. KL 19.691242218017578
    Iteration 6300. Loss 106.8331527709961. KL 19.24445343017578
    Iteration 6400. Loss 104.03341674804688. KL 19.696138381958008
    Iteration 6500. Loss 98.54393005371094. KL 19.088619232177734
    Iteration 6600. Loss 105.12592315673828. KL 19.1871337890625
    Iteration 6700. Loss 105.50454711914062. KL 19.622892379760742
    Iteration 6800. Loss 104.51860809326172. KL 18.811485290527344
    Iteration 6900. Loss 102.87434387207031. KL 19.169269561767578
    Iteration 7000. Loss 100.8294677734375. KL 19.846006393432617
    Iteration 7100. Loss 101.88845825195312. KL 18.935867309570312
    Iteration 7200. Loss 107.11431884765625. KL 18.98651123046875
    Iteration 7300. Loss 109.47294616699219. KL 19.818172454833984
    Iteration 7400. Loss 101.6451416015625. KL 19.717487335205078
    Iteration 7500. Loss 106.61439514160156. KL 19.449527740478516
    Iteration 7600. Loss 101.13935852050781. KL 18.896699905395508
    Iteration 7700. Loss 102.81201171875. KL 19.145170211791992
    Iteration 7800. Loss 102.26914978027344. KL 19.4200439453125
    Iteration 7900. Loss 100.77509307861328. KL 19.934852600097656
    Iteration 8000. Loss 105.97528839111328. KL 19.732383728027344
    Iteration 8100. Loss 99.77294921875. KL 19.025257110595703
    Iteration 8200. Loss 104.75892639160156. KL 19.709518432617188
    Iteration 8300. Loss 104.98646545410156. KL 19.123580932617188
    Iteration 8400. Loss 102.78396606445312. KL 18.542850494384766
    Iteration 8500. Loss 98.69320678710938. KL 18.84295654296875
    Iteration 8600. Loss 97.60421752929688. KL 18.89189910888672
    Iteration 8700. Loss 100.128662109375. KL 18.16009521484375
    Iteration 8800. Loss 100.00633239746094. KL 19.163223266601562
    Iteration 8900. Loss 100.12950134277344. KL 19.470077514648438
    Iteration 9000. Loss 104.17083740234375. KL 19.365875244140625
    Iteration 9100. Loss 102.51815795898438. KL 19.76079559326172
    Iteration 9200. Loss 98.08135986328125. KL 19.184261322021484
    Iteration 9300. Loss 105.66838073730469. KL 19.566570281982422
    Iteration 9400. Loss 94.75746154785156. KL 18.388275146484375
    Iteration 9500. Loss 100.62859344482422. KL 18.82201385498047
    Iteration 9600. Loss 103.2626724243164. KL 20.16010093688965
    Iteration 9700. Loss 102.6614990234375. KL 19.639568328857422
    Iteration 9800. Loss 94.9361343383789. KL 19.19310760498047
    Iteration 9900. Loss 101.02583312988281. KL 19.19292449951172

### Visualize the loss in time

```python
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

### Generate samples and latent interpolations

> Indented block

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about amortized variational inference

* What do you observe about sample quality and reconstruction quality?
* What do you observe about the ELBO and KL term?
* Which gradient estimation method is used to compute the gradients with respect to the encoder parameters?

# KL annealing

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)}  - \alpha \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

Where $\alpha$ changes during training, to weigh in the KL term more.

```python
tf.reset_default_graph()
```

```python
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

```python
kl_coefficient = tf.Variable(
    tf.zeros(shape=(1), dtype=tf.float32),
    trainable=False,
    name='kl_coeff')

kl_step = 1.0 / TRAINING_STEPS
update_kl_coeff = tf.assign(kl_coefficient, kl_coefficient + kl_step)
```

```python
variational_posterior = encoder(real_data)# Maximize the data likelihodd and minimize the KL divergence between the prior
# and posterior. We use the exact same loss as in the SVI case.
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
train_elbo = likelihood_term - kl_coefficient * kl_term
loss =  - train_elbo
```

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
variables_update_op = optimizer.minimize(loss, var_list=autoencoder_vars)

# Ensure that a variable update is followed by an update in the KL coefficient.
with tf.control_dependencies([variables_update_op]):
  update_op = tf.identity(update_kl_coeff)
```

## Training

```python
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

    W0715 15:52:22.507540 4672783808 deprecation.py:323] From /anaconda3/lib/python3.7/site-packages/tensorflow/python/util/tf_should_use.py:193: initialize_all_variables (from tensorflow.python.ops.variables) is deprecated and will be removed after 2017-03-02.
    Instructions for updating:
    Use `tf.global_variables_initializer` instead.

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

    Iteration 0. Loss [540.14087]. KL 0.04890408366918564
    Iteration 100. Loss [199.93692]. KL 70.31986236572266
    Iteration 200. Loss [158.17885]. KL 68.29264068603516
    Iteration 300. Loss [119.23048]. KL 59.43687438964844
    Iteration 400. Loss [104.75813]. KL 56.192832946777344
    Iteration 500. Loss [105.39626]. KL 49.35320281982422
    Iteration 600. Loss [97.36128]. KL 49.204437255859375
    Iteration 700. Loss [94.2336]. KL 40.21452713012695
    Iteration 800. Loss [95.09839]. KL 35.83056640625
    Iteration 900. Loss [91.84171]. KL 35.596405029296875
    Iteration 1000. Loss [99.84914]. KL 35.6655158996582
    Iteration 1100. Loss [102.530396]. KL 32.10776138305664
    Iteration 1200. Loss [91.94745]. KL 34.0959587097168
    Iteration 1300. Loss [101.78357]. KL 31.88787841796875
    Iteration 1400. Loss [89.15325]. KL 30.364116668701172
    Iteration 1500. Loss [94.65003]. KL 31.283523559570312
    Iteration 1600. Loss [96.22745]. KL 29.73382568359375
    Iteration 1700. Loss [89.04879]. KL 29.58684539794922
    Iteration 1800. Loss [93.3224]. KL 29.75642204284668
    Iteration 1900. Loss [89.121376]. KL 29.29058074951172
    Iteration 2000. Loss [86.97606]. KL 28.70804214477539
    Iteration 2100. Loss [81.24902]. KL 28.902484893798828
    Iteration 2200. Loss [83.018036]. KL 28.541650772094727
    Iteration 2300. Loss [91.816055]. KL 28.694034576416016
    Iteration 2400. Loss [91.886894]. KL 28.517147064208984
    Iteration 2500. Loss [91.91321]. KL 28.265640258789062
    Iteration 2600. Loss [84.27647]. KL 27.58275032043457
    Iteration 2700. Loss [90.052826]. KL 27.69062042236328
    Iteration 2800. Loss [94.781586]. KL 27.390445709228516
    Iteration 2900. Loss [91.79245]. KL 27.005218505859375
    Iteration 3000. Loss [92.32571]. KL 25.370742797851562
    Iteration 3100. Loss [88.81453]. KL 26.762733459472656
    Iteration 3200. Loss [90.82612]. KL 26.097944259643555
    Iteration 3300. Loss [90.46532]. KL 24.50055694580078
    Iteration 3400. Loss [93.961075]. KL 26.91839599609375
    Iteration 3500. Loss [85.490715]. KL 25.969730377197266
    Iteration 3600. Loss [91.61385]. KL 26.85202407836914
    Iteration 3700. Loss [96.483536]. KL 25.482364654541016
    Iteration 3800. Loss [94.82649]. KL 26.293773651123047
    Iteration 3900. Loss [97.695946]. KL 25.082019805908203
    Iteration 4000. Loss [94.28689]. KL 24.24665069580078
    Iteration 4100. Loss [87.30277]. KL 24.85059356689453
    Iteration 4200. Loss [90.992035]. KL 23.834518432617188
    Iteration 4300. Loss [95.25837]. KL 25.463592529296875
    Iteration 4400. Loss [102.102936]. KL 24.69823455810547
    Iteration 4500. Loss [91.89266]. KL 24.80847930908203
    Iteration 4600. Loss [95.95828]. KL 24.694360733032227
    Iteration 4700. Loss [94.15628]. KL 24.71017837524414
    Iteration 4800. Loss [101.23945]. KL 24.380550384521484
    Iteration 4900. Loss [96.551674]. KL 23.871843338012695
    Iteration 5000. Loss [97.09316]. KL 23.649051666259766
    Iteration 5100. Loss [90.89116]. KL 24.07223892211914
    Iteration 5200. Loss [96.468315]. KL 25.30548095703125
    Iteration 5300. Loss [88.045265]. KL 23.267858505249023
    Iteration 5400. Loss [93.23415]. KL 23.363754272460938
    Iteration 5500. Loss [100.58287]. KL 24.101051330566406
    Iteration 5600. Loss [92.58667]. KL 23.428377151489258
    Iteration 5700. Loss [95.1404]. KL 22.430952072143555
    Iteration 5800. Loss [98.40685]. KL 22.85318374633789
    Iteration 5900. Loss [92.41812]. KL 22.730545043945312
    Iteration 6000. Loss [93.50008]. KL 22.5500431060791
    Iteration 6100. Loss [94.336945]. KL 22.30872917175293
    Iteration 6200. Loss [99.464745]. KL 22.114582061767578
    Iteration 6300. Loss [100.76575]. KL 23.462177276611328
    Iteration 6400. Loss [95.75052]. KL 22.453447341918945
    Iteration 6500. Loss [89.122955]. KL 21.759613037109375
    Iteration 6600. Loss [104.24829]. KL 21.740550994873047
    Iteration 6700. Loss [92.29272]. KL 22.474525451660156
    Iteration 6800. Loss [94.40081]. KL 22.50080108642578
    Iteration 6900. Loss [102.171326]. KL 22.135616302490234
    Iteration 7000. Loss [95.10802]. KL 22.734222412109375
    Iteration 7100. Loss [102.47019]. KL 22.65586280822754
    Iteration 7200. Loss [90.98844]. KL 20.730186462402344
    Iteration 7300. Loss [93.12424]. KL 21.640308380126953
    Iteration 7400. Loss [101.63361]. KL 21.304807662963867
    Iteration 7500. Loss [98.62222]. KL 21.820524215698242
    Iteration 7600. Loss [100.14249]. KL 21.75533676147461
    Iteration 7700. Loss [102.1537]. KL 21.460617065429688
    Iteration 7800. Loss [95.63809]. KL 20.831697463989258
    Iteration 7900. Loss [107.83273]. KL 21.953369140625
    Iteration 8000. Loss [98.962494]. KL 21.2552490234375
    Iteration 8100. Loss [96.57939]. KL 20.178586959838867
    Iteration 8200. Loss [98.8622]. KL 20.882596969604492
    Iteration 8300. Loss [95.91652]. KL 21.824268341064453
    Iteration 8400. Loss [101.63084]. KL 21.216724395751953
    Iteration 8500. Loss [94.91183]. KL 20.33702278137207
    Iteration 8600. Loss [102.232834]. KL 20.5839786529541
    Iteration 8700. Loss [104.95986]. KL 20.58934211730957
    Iteration 8800. Loss [100.7558]. KL 21.401830673217773
    Iteration 8900. Loss [98.60055]. KL 20.360279083251953
    Iteration 9000. Loss [96.82796]. KL 19.279766082763672
    Iteration 9100. Loss [103.07653]. KL 20.187301635742188
    Iteration 9200. Loss [103.05805]. KL 20.406352996826172
    Iteration 9300. Loss [97.76108]. KL 19.578073501586914
    Iteration 9400. Loss [105.954414]. KL 21.17473602294922
    Iteration 9500. Loss [97.83456]. KL 20.006214141845703
    Iteration 9600. Loss [99.97221]. KL 20.190013885498047
    Iteration 9700. Loss [99.3232]. KL 20.683671951293945
    Iteration 9800. Loss [102.7419]. KL 20.11018180847168
    Iteration 9900. Loss [98.26511]. KL 20.12986183166504

```python
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about KL annealing

* What do you observe about the KL behaviour throughout training as opposed to amortized variational inference without any KL annealing?

# Constrained optimization

Instead of using KL annealing, constrained optimization can be used to automatically tuned the relative weight of the likelihood and kl terms. This removes the need to manually create an optimization schedule, which can be problem specific.

The objective now becomes:

\begin{equation}
  \text{minimize } \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) \text{ such that }  \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)} > \alpha
\end{equation}

This can be solved using the use of Lagrange multipliers. The objective becomes:

\begin{equation}
  \text{minimize }  \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) + \lambda  (\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z)))
\end{equation}

The difference compared to the KL annealing is that:

   * $\lambda$ is a learned parameter - it will be learned using stochastic gradient descent, like the network parameters.  The difference is that the lagrangian has to solve a maximization problem. You can see this intuitively: the graadient with respect to $\lambda$ in the objective above is $\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))$. If $ \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))> 0$, the constraint is not being satisfied, so the value of the lagrangian needs to increase. This will be done by doing gradient ascent, instead of gradient descent. Note that for $\lambda$ to be a valid lagranian in a minimization problem, it has to be positive.
   * The practicioner has to specify the hyperparameter $\alpha$, which determines the reoncstruction quality of the model.
   * the coefficient is in front of the likelihood term, not the KL term. This is mainly for convenience, as it is easier to specify the hyperparameter $\alpha$ for the likelihood (reconstruction loss).

For more assumptions made by this method, see the Karush–Kuhn–Tucker conditions.

For more information, see:
  * http://bayesiandeeplearning.org/2018/papers/33.pdf

```python
tf.reset_default_graph()
```

```python
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

```python
# Unlike in the KL annealing case, we learn the coefficient.
# We initialize the lagrangian such that after the softplus the coefficient is
# approximately 1.
lagrangian_var = tf.Variable(
    0.55 * tf.ones(shape=(1), dtype=tf.float32),
    trainable=True,
    name='lagrangian')

# Ensure that the lagrangian is positive and has stable dynamics.
lagrangian = tf.nn.softplus(lagrangian_var)
```

```python
# How good do we want the reconstruction loss to be?
# We can look at previous runs to get an idea what a reasonable value would be.
reconstruction_target = - 90
```

```python
variational_posterior = encoder(real_data)
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
# Note: now the elbo is a different quanitity to what we optimize.
train_elbo = likelihood_term - kl_term
loss =  kl_term + lagrangian * (reconstruction_target - likelihood_term )
```

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
tf.trainable_variables()
```

    [<tf.Variable 'lagrangian:0' shape=(1,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d/kernel:0' shape=(5, 5, 1, 8) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d/bias:0' shape=(8,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_1/kernel:0' shape=(5, 5, 8, 16) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_1/bias:0' shape=(16,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_2/kernel:0' shape=(5, 5, 16, 32) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_2/bias:0' shape=(32,) dtype=float32_ref>,
     <tf.Variable 'encoder/dense/kernel:0' shape=(1568, 10) dtype=float32_ref>,
     <tf.Variable 'encoder/dense/bias:0' shape=(10,) dtype=float32_ref>,
     <tf.Variable 'encoder/dense_1/kernel:0' shape=(1568, 10) dtype=float32_ref>,
     <tf.Variable 'encoder/dense_1/bias:0' shape=(10,) dtype=float32_ref>,
     <tf.Variable 'decoder/dense/kernel:0' shape=(10, 3136) dtype=float32_ref>,
     <tf.Variable 'decoder/dense/bias:0' shape=(3136,) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose/kernel:0' shape=(5, 5, 32, 64) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose/bias:0' shape=(32,) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose_1/kernel:0' shape=(5, 5, 1, 32) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose_1/bias:0' shape=(1,) dtype=float32_ref>]

```python
lagrangian_optimizer = tf.train.GradientDescentOptimizer(0.001)

# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
autoencoder_variables_update_op = optimizer.minimize(
      loss, var_list=autoencoder_vars)

# Ensure that a variable update is followed by an update to the Lagrangian.
with tf.control_dependencies([autoencoder_variables_update_op]):
  # Ensure that the lagrangian solves a maximization problem instead of a
  # minimization problem by changing the sign of the loss function.
  update_op = lagrangian_optimizer.minimize(- loss, var_list=[lagrangian_var])
```

## Training

```python
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []
lagrangian_values = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl, lag_val  = sess.run(
          [loss, likelihood_term, kl_term, lagrangian])
    print('Iteration {}. Loss {}. KL {}. Lagrangian {}'.format(
          i, iteration_loss, iteration_kl, lag_val))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
    lagrangian_values.append(lag_val)
```

    Iteration 0. Loss [538.96967]. KL 0.00980534590780735. Lagrangian [1.1973238]
    Iteration 100. Loss [2012.2874]. KL 26.05329132080078. Lagrangian [18.439686]
    Iteration 200. Loss [2037.9198]. KL 43.41221237182617. Lagrangian [27.876625]
    Iteration 300. Loss [1312.1792]. KL 53.85942077636719. Lagrangian [33.11151]
    Iteration 400. Loss [604.27075]. KL 59.25827407836914. Lagrangian [35.38244]
    Iteration 500. Loss [103.39664]. KL 64.09086608886719. Lagrangian [36.683895]
    Iteration 600. Loss [321.0326]. KL 51.04595184326172. Lagrangian [37.5285]
    Iteration 700. Loss [611.9413]. KL 52.78070831298828. Lagrangian [38.12265]
    Iteration 800. Loss [193.57184]. KL 54.54662322998047. Lagrangian [38.41363]
    Iteration 900. Loss [509.35577]. KL 49.528629302978516. Lagrangian [38.68292]
    Iteration 1000. Loss [155.02695]. KL 50.559043884277344. Lagrangian [38.745945]
    Iteration 1100. Loss [319.78564]. KL 49.118499755859375. Lagrangian [38.680435]
    Iteration 1200. Loss [-131.5185]. KL 48.60651397705078. Lagrangian [38.55569]
    Iteration 1300. Loss [157.66992]. KL 47.537288665771484. Lagrangian [38.416492]
    Iteration 1400. Loss [256.60852]. KL 45.675880432128906. Lagrangian [38.176632]
    Iteration 1500. Loss [-183.12325]. KL 47.668792724609375. Lagrangian [37.841064]
    Iteration 1600. Loss [-26.41735]. KL 47.68638610839844. Lagrangian [37.49614]
    Iteration 1700. Loss [-214.85039]. KL 46.65449523925781. Lagrangian [37.068287]
    Iteration 1800. Loss [-311.81882]. KL 46.79855728149414. Lagrangian [36.53549]
    Iteration 1900. Loss [-253.39359]. KL 45.166725158691406. Lagrangian [36.012287]
    Iteration 2000. Loss [-20.606178]. KL 45.8177490234375. Lagrangian [35.481823]
    Iteration 2100. Loss [-380.71466]. KL 43.924102783203125. Lagrangian [34.865334]
    Iteration 2200. Loss [-239.90263]. KL 45.6566047668457. Lagrangian [34.282005]
    Iteration 2300. Loss [-88.56151]. KL 43.67646789550781. Lagrangian [33.63299]
    Iteration 2400. Loss [-106.883644]. KL 43.311317443847656. Lagrangian [32.96393]
    Iteration 2500. Loss [-404.17126]. KL 43.54340362548828. Lagrangian [32.240883]
    Iteration 2600. Loss [-346.92252]. KL 43.05750274658203. Lagrangian [31.482908]
    Iteration 2700. Loss [-71.27568]. KL 42.65627670288086. Lagrangian [30.70962]
    Iteration 2800. Loss [-223.26251]. KL 42.003822326660156. Lagrangian [29.930614]
    Iteration 2900. Loss [-258.41284]. KL 40.27805709838867. Lagrangian [29.093626]
    Iteration 3000. Loss [-134.62718]. KL 43.185150146484375. Lagrangian [28.308224]
    Iteration 3100. Loss [-193.20137]. KL 40.280181884765625. Lagrangian [27.477444]
    Iteration 3200. Loss [-105.128006]. KL 41.28340148925781. Lagrangian [26.587652]
    Iteration 3300. Loss [-309.9415]. KL 40.566864013671875. Lagrangian [25.71214]
    Iteration 3400. Loss [-145.85098]. KL 41.6849365234375. Lagrangian [24.774744]
    Iteration 3500. Loss [-32.124664]. KL 40.62626647949219. Lagrangian [23.826397]
    Iteration 3600. Loss [-156.75305]. KL 39.939598083496094. Lagrangian [22.842669]
    Iteration 3700. Loss [-234.57065]. KL 38.174659729003906. Lagrangian [21.85806]
    Iteration 3800. Loss [-181.11844]. KL 38.56774139404297. Lagrangian [20.88989]
    Iteration 3900. Loss [-259.78387]. KL 38.60774230957031. Lagrangian [19.920109]
    Iteration 4000. Loss [-31.917583]. KL 38.87125778198242. Lagrangian [19.00201]
    Iteration 4100. Loss [-121.71157]. KL 37.418479919433594. Lagrangian [18.074179]
    Iteration 4200. Loss [-162.5707]. KL 40.08145523071289. Lagrangian [17.069633]
    Iteration 4300. Loss [-150.18365]. KL 36.11170196533203. Lagrangian [15.975876]
    Iteration 4400. Loss [-147.69156]. KL 36.65587615966797. Lagrangian [14.915364]
    Iteration 4500. Loss [-181.67966]. KL 35.739967346191406. Lagrangian [13.911487]
    Iteration 4600. Loss [-96.17489]. KL 35.353233337402344. Lagrangian [12.903025]
    Iteration 4700. Loss [-115.332184]. KL 34.46046829223633. Lagrangian [11.859877]
    Iteration 4800. Loss [-81.467514]. KL 33.849369049072266. Lagrangian [10.780733]
    Iteration 4900. Loss [-83.52796]. KL 34.126991271972656. Lagrangian [9.708743]
    Iteration 5000. Loss [-28.11337]. KL 31.60072135925293. Lagrangian [8.645763]
    Iteration 5100. Loss [-60.07852]. KL 32.564422607421875. Lagrangian [7.5990987]
    Iteration 5200. Loss [-28.396912]. KL 31.27129364013672. Lagrangian [6.557491]
    Iteration 5300. Loss [-41.119373]. KL 30.55807876586914. Lagrangian [5.4892154]
    Iteration 5400. Loss [-28.609264]. KL 29.060646057128906. Lagrangian [4.4663324]
    Iteration 5500. Loss [-9.711636]. KL 27.815330505371094. Lagrangian [3.4340072]
    Iteration 5600. Loss [-8.41758]. KL 25.981338500976562. Lagrangian [2.5337074]
    Iteration 5700. Loss [11.221826]. KL 23.470190048217773. Lagrangian [1.8373184]
    Iteration 5800. Loss [13.211808]. KL 21.701915740966797. Lagrangian [1.322103]
    Iteration 5900. Loss [12.488285]. KL 20.661758422851562. Lagrangian [1.0462322]
    Iteration 6000. Loss [9.482178]. KL 19.263051986694336. Lagrangian [0.82997036]
    Iteration 6100. Loss [14.167581]. KL 18.332239151000977. Lagrangian [0.7134634]
    Iteration 6200. Loss [16.62278]. KL 18.950946807861328. Lagrangian [0.646816]
    Iteration 6300. Loss [16.947508]. KL 17.118423461914062. Lagrangian [0.60027325]
    Iteration 6400. Loss [15.33502]. KL 16.43392562866211. Lagrangian [0.56152964]
    Iteration 6500. Loss [15.829773]. KL 15.905801773071289. Lagrangian [0.552094]
    Iteration 6600. Loss [14.940361]. KL 16.844772338867188. Lagrangian [0.5408143]
    Iteration 6700. Loss [14.527574]. KL 15.577510833740234. Lagrangian [0.53791887]
    Iteration 6800. Loss [16.135342]. KL 16.559040069580078. Lagrangian [0.5398146]
    Iteration 6900. Loss [19.023605]. KL 16.628189086914062. Lagrangian [0.5391308]
    Iteration 7000. Loss [12.499035]. KL 16.135868072509766. Lagrangian [0.5362083]
    Iteration 7100. Loss [15.16113]. KL 15.600274085998535. Lagrangian [0.52921474]
    Iteration 7200. Loss [15.442754]. KL 16.5613956451416. Lagrangian [0.52967167]
    Iteration 7300. Loss [17.04506]. KL 16.015830993652344. Lagrangian [0.5302499]
    Iteration 7400. Loss [16.218252]. KL 16.51810073852539. Lagrangian [0.52346474]
    Iteration 7500. Loss [15.971554]. KL 15.572644233703613. Lagrangian [0.52481085]
    Iteration 7600. Loss [14.602805]. KL 16.155397415161133. Lagrangian [0.52090895]
    Iteration 7700. Loss [10.875221]. KL 15.068178176879883. Lagrangian [0.513955]
    Iteration 7800. Loss [16.052803]. KL 15.61003303527832. Lagrangian [0.5093886]
    Iteration 7900. Loss [18.802227]. KL 15.69420051574707. Lagrangian [0.5091186]
    Iteration 8000. Loss [13.41296]. KL 15.310225486755371. Lagrangian [0.50918627]
    Iteration 8100. Loss [13.967747]. KL 15.255514144897461. Lagrangian [0.51302797]
    Iteration 8200. Loss [15.583532]. KL 16.596294403076172. Lagrangian [0.5147781]
    Iteration 8300. Loss [18.125237]. KL 15.539207458496094. Lagrangian [0.5153093]
    Iteration 8400. Loss [16.278257]. KL 16.647384643554688. Lagrangian [0.51031715]
    Iteration 8500. Loss [16.442034]. KL 15.341384887695312. Lagrangian [0.5090413]
    Iteration 8600. Loss [16.717308]. KL 15.561285018920898. Lagrangian [0.5088723]
    Iteration 8700. Loss [16.387648]. KL 16.300559997558594. Lagrangian [0.512679]
    Iteration 8800. Loss [17.292738]. KL 14.986751556396484. Lagrangian [0.5063022]
    Iteration 8900. Loss [14.255327]. KL 14.89797592163086. Lagrangian [0.501947]
    Iteration 9000. Loss [16.814037]. KL 16.098064422607422. Lagrangian [0.5018876]
    Iteration 9100. Loss [12.317408]. KL 14.925344467163086. Lagrangian [0.4945041]
    Iteration 9200. Loss [14.718001]. KL 14.65495491027832. Lagrangian [0.49903888]
    Iteration 9300. Loss [14.590738]. KL 15.930028915405273. Lagrangian [0.49838033]
    Iteration 9400. Loss [19.41078]. KL 15.671988487243652. Lagrangian [0.50269675]
    Iteration 9500. Loss [15.791254]. KL 15.040611267089844. Lagrangian [0.5015253]
    Iteration 9600. Loss [13.987355]. KL 15.799518585205078. Lagrangian [0.49544224]
    Iteration 9700. Loss [16.478256]. KL 14.903437614440918. Lagrangian [0.49357888]
    Iteration 9800. Loss [13.422607]. KL 15.213420867919922. Lagrangian [0.4965107]
    Iteration 9900. Loss [16.314394]. KL 15.002989768981934. Lagrangian [0.49242073]

```python
fig, axes = plt.subplots(2, 2, figsize=(2*8, 2* 5))

axes[0, 0].plot(losses, label='Negative ELBO')
axes[0, 0].set_title('Time', fontsize=15)
axes[0, 0].legend()

axes[0, 1].plot(kls, label='KL')
axes[0, 1].set_title('Time', fontsize=15)
axes[0, 1].legend()

axes[1, 0].plot(likelihood_terms, label='Likelihood Term')
axes[1, 0].set_title('Time', fontsize=15)
axes[1, 0].legend()

axes[1, 1].plot(lagrangian_values, label='Lagrangian Values')
axes[1, 1].set_title('Time', fontsize=15)
axes[1, 1].legend()
```

### Generate samples and latent interpolations

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Discussion about constrained optimization

* What do you observe about the behaviour of the likelihood and KL term throughout training? How is it different than in Amortized Variational inference with and without KL annealing?
* What do you notice about the behaviour of the lagrangian during training?
* How would you choose the hyperparameters of the constrained optimization problem?

Machine Learning Summer School 2019 (London, UK)/tutorials/vi/MLSS2019_VAEs_solution.ipynb