### Checklist for submission

It is extremely important to make sure that:

1. Everything runs as expected (no bugs when running cells);
2. The output from each cell corresponds to its code (don't change any cell's contents without rerunning it afterwards);
3. All outputs are present (don't delete any of the outputs);
4. Fill in all the places that say `# YOUR CODE HERE`, or "**Your answer:** (fill in here)".
5. Never copy/paste any notebook cells. Inserting new cells is allowed, but it should not be necessary.
6. The notebook contains some hidden metadata which is important during our grading process. **Make sure not to corrupt any of this metadata!** The metadata may for example be corrupted if you copy/paste any notebook cells, or if you perform an unsuccessful git merge / git pull. It may also be pruned completely if using Google Colab, so watch out for this. Searching for "nbgrader" when opening the notebook in a text editor should take you to the important metadata entries.
7. Although we will try our very best to avoid this, it may happen that bugs are found after an assignment is released, and that we will push an updated version of the assignment to GitHub. If this happens, it is important that you update to the new version, while making sure the notebook metadata is properly updated as well. The safest way to make sure nothing gets messed up is to start from scratch on a clean updated version of the notebook, copy/pasting your code from the cells of the previous version into the cells of the new version.
8. If you need to have multiple parallel versions of this notebook, make sure not to move them to another directory.
9. Although not forced to work exclusively in the course `conda` environment, you need to make sure that the notebook will run in that environment, i.e. that you have not added any additional dependencies.

**FOR HA1, HA2, HA3 ONLY:** Failing to meet any of these requirements might lead to either a subtraction of POEs (at best) or a request for resubmission (at worst).

We advise you to perform the following steps before submission to ensure that requirements 1, 2, and 3 are always met: **Restart the kernel** (in the menubar, select Kernel$\rightarrow$Restart) and then **run all cells** (in the menubar, select Cell$\rightarrow$Run All). This might require a bit of time, so plan ahead for this (and possibly use Google Cloud's GPU in HA1 and HA2 for this step). Finally press the "Save and Checkout" button before handing in, to make sure that all your changes are saved to this .ipynb file.

### Fill in name of notebook file
This might seem silly, but the version check below needs to know the filename of the current notebook, which is not trivial to find out programmatically.

You might want to have several parallel versions of the notebook, and it is fine to rename the notebook as long as it stays in the same directory. **However**, if you do rename it, you also need to update its own filename below:

```python
# nb_fname = "XXX.ipynb"
```

### Fill in group number and member names (use NAME2 and GROUP only for HA1, HA2 and HA3):

```python
NAME1 = ""
NAME2 = ""
GROUP = ""
```

### Check Python version

```python
from platform import python_version_tuple
assert python_version_tuple()[:2] == ('3','7'), "You are not running Python 3.7. Make sure to run Python through the course Conda environment."
```

### Check that notebook server has access to all required resources, and that notebook has not moved

```python
import os
nb_dirname = os.path.abspath('')
assignment_name = os.path.basename(nb_dirname)
assert assignment_name in ['IHA1', 'IHA2', 'HA1', 'HA2', 'HA3'], \
    '[ERROR] The notebook appears to have been moved from its original directory'
```

### Verify correct nb_fname

```python
from IPython.display import display, HTML
try:
    display(HTML(r''.format(nb_fname=nb_fname)))
except NameError:
    assert False, 'Make sure to fill in the nb_fname variable above!'
```

### Verify that your notebook is up-to-date and not corrupted in any way

```python
import sys
sys.path.append('..')
from ha_utils import check_notebook_uptodate_and_not_corrupted
check_notebook_uptodate_and_not_corrupted(nb_dirname, nb_fname)
```

# HA2:  Part 2 - Transformers and self-attention
$$
\renewcommand{\vec}[1]{#1}
\def\x{\vec{x}}
\def\y{\vec{y}}
\def\dim{d}
\def\w{W}
\def\wu{Z}
\def\R{\mathbb{R}}
\def\linMap{W}
% Query, key and val
\def\q{\vec{q}}
\def\k{\vec{k}}
\def\v{\vec{v}}
\def\Wq{\linMap_Q}
\def\Wk{\linMap_K}
\def\Wv{\linMap_V}
$$
*You should have completed part 1 before starting with this one*

In this part we will take a closer look at the transformer architecture and the self-attention operation.
We will start with basic self-attention and gradually construct an actual self-attention module.
Finally we will construct a complete transformer and test it on an actual problem.

The focus is on a conceptual understanding of the transformer but you will have to implement a few key elements of a transformer. Along the way we will try to give some best practices for constructing a more complex network architecture.

Let's start with importing the module's we are going to need:

```python
import torch
import torch.nn as torch_nn
import torch.nn.functional as F
```

# 1. Basic self-attention

The key-stone of the transformer architecture, self-attention is a sequence-sequence operation which transforms a sequence of input vectors $\x_1, \dots \x_t$ to output vectors $\y_1, \dots \y_t$.
Remember that all vectors have the same dimension $\dim$, i.e. $\x_i, \y_i \in \R^{\dim}, \forall i = 1, \dots t$.

## Weighted average
The actual transformation is a simple weighted average
$$
\y_i = \sum_{j} \x_j \w_{ji}.
$$

In an actual transformer, weighted averages are computed often and for long sequences. Therefore, the implementation must be fast in order for training to be even possible.
With high-level frameworks such as `pytorch`, the key to fast code is often to reduce loops and instead express computations as matrix operations.

**[3 points]** Complete the function snippet below to implement simple weight sharing.

To pass this part of the assignment your implementation only has to be correct, not efficient, but to get one extra point, you must implement it with just a single for loop. For two extra points, do it without any loops at all.

*Hint*: Take a look at how `torch.bmm` is used later in the implementation

```python
def weighted_avg(x, weights):
    """Weighted average
    Calculates a weighted average of a batch of sequences of vectors.

    Args:
        x (torch.Tensor): Shape (batch_size, dim, seq_len)
        weights (torch.Tensor): Shape (batch_size, seq_len, seq_len)

    Returns:
        y (torch.Tensor): Shape (batch_size, dim, seq_len)
    """
    # YOUR CODE HERE
```

Make sure to test your implementation with the unit tests below.
The tests cover:

1. Dimensionality
2. Uniform weights $\w_{ji} = \frac{1}{t}$ should produce $y_i: y_i = \frac{1}{\dim} \sum_{j} x_j,\, \forall i = 1, \dots t$
 (i.e., every $y_i$ is an average of the input sequence).
3. A specific numerical example with batch size = 2, $t = 2,\, \dim=1$.

```python
def test_weighted_avg(function):
    """
    Args:
        function: Implementation to test
    """
    # Testing dimension of averaged tensor.
    batch_size, dim, seq_len = 5, 2, 3
    x = torch.rand(batch_size, dim, seq_len)
    weights = torch.rand(batch_size, seq_len, seq_len)
    y = function(x, weights)
    assert y.shape == (batch_size, dim, seq_len), "Dimension error: expected y to have shape {}, got {}.".format(
        (batch_size, seq_len, dim), tuple(y.shape))

    # Testing uniform weights preserve x.
    batch_size, dim, seq_len = 5, 2, 3
    x = torch.rand(batch_size, dim, seq_len)
    weights = torch.ones((batch_size, seq_len, seq_len)).float() / seq_len
    y = function(x, weights)
    assert all(torch.allclose(y_b.mean(1), y_b[:, 0]) for y_b in y),\
        "Numerical error: With uniform weights, expected y_i = y_j forall i, j (within each batch)."
    assert all(torch.allclose(y_b.mean(1), x_b.mean(1)) for (x_b, y_b) in zip(x, y)),\
        "Numerical error: With uniform weights, expected y_i = sum_j x_j, for all i"

    # Actual numerical example.
    x = torch.tensor([4, 1]).reshape((1, 1, 2)).float()
    unnorm_weights = torch.arange(1, 5).reshape((1, 2, 2)).float()
    scale = unnorm_weights.sum(1).reshape((1, 1, 2))
    weights = unnorm_weights / scale

    y = function(x, weights)
    y_true = torch.tensor([7/4, 2]).reshape(1, 1, 2).float()
    assert torch.allclose(y, y_true), "Numerical error, expected: {}, got {}".format(y_true, y)

    print("Test passed.")

test_weighted_avg(function=weighted_avg)
```

## Defining weights through the dot product
A simple way to define $\w_{ji}$ is with the dot product

$$
\wu_{ji} = \x_j^T \x_i.
$$
which maps the pair of input vectors to a non-negative scalar, $\R^{\dim \times \dim} \to [0, \infty)$.
We then use a softmax to obtain normalised $\w_{ji} \in (0, 1]$:

$$
\w_{ji} = \frac{ e^{\wu_{ji}} }{ \sum_j e^{\wu_{ji}} }.
$$

**[2 points]** What is the difference between these weights and the weights in ordinary networks, e.g. a CNN?

**Your answer:** (fill in here)

The dot product is essential for calculating the weights. As we progress, we will make slight modifications to the inputs but we will still base it around a function which calculates a softmax-normalized dot product. Therefore, you need to complete the implementation below **[3 points]**:

Again, this function will be evaluated often and for long sequences in the transformer block. For all points, implement it without using for loops.

```python
def normalized_dot_product(v_1, v_2):
    """Normalized dot products between all pairs of vectors in a sequence
    Takes two batches of sequences of vectors as input.
    Sequences in the batch are processed independently.
    The normalization is done with a softmax function along the columns of the weight matrices.

    Args:
        v_1 (torch.Tensor): Shape (batch_size, dim, seq_len)
        v_2 (torch.Tensor): Shape (batch_size, dim, seq_len)

    Returns:
        norm_dot_prod (torch.Tensor): Shape (batch_size, seq_len, seq_len)
    """

    # YOUR CODE HERE
```

Make sure to test your implementation with the unit tests below.
The tests cover:

1. Dimensionality
2. Normalized in the correct dimension
3. A specific numerical example

```python
import numpy as np

def test_normalized_dot_product(function):
    """
    Args:
        function: Implementation to test
    """

    batch_size, dim, seq_len = 5, 2, 3
    v_1 = torch.rand(batch_size, dim, seq_len)
    v_2 = torch.rand(batch_size, dim, seq_len)
    weights = function(v_1, v_2)

    # Testing dimension of weights.
    assert weights.shape == (batch_size, seq_len, seq_len),\
    "Dimension error: expected weights to have shape {}, got {}.".format(
        (batch_size, seq_len, seq_len), tuple(weights.shape))

    # Testing weights non-negative
    # (Boolean tensor's can be reduced to a single boolean)
    assert not (weights < 0.0).any() ,\
    "Value error: expected weights to be non-negative."

    # Testing weights smaller than one
    assert (weights < 1.0).all() ,\
    "Value error: expected weights to be non-negative."

    assert torch.allclose(weights.sum(1), torch.ones((batch_size, seq_len))),\
        "ValueError: expected columns (dim 1) to sum to 1.0"

    # Actual numerical example
    v_1 = torch.tensor([[1, 2], [-1, 1]]).float().reshape((1, 2, 2))
    v_2 = torch.tensor([[1, 0], [1, -1]]).float().reshape((1, 2, 2))
    e = np.exp(1)
    true_weights = torch.tensor([
        [1 / (e**3 + 1), e**2 / (e**2 + 1)],
        [e**3 / (e**3 + 1), 1 / (e**2 + 1)]
    ]).reshape((1, 2, 2))
    weights =  function(v_1, v_2)
    assert torch.allclose(true_weights, weights),\
    "Numerical error: expected {}, got {}.".format(true_weights, weights)

    print("Test passed.")

test_normalized_dot_product(function=normalized_dot_product)
```

That's it, we have now the building blocks needed for basic self-attention:

```python
def basic_self_attention(x):
    """Basic self-attention
    Transforms a batch of sequences of vectors.

    Args:
        x (torch.Tensor): Shape (batch_size, dim, seq_len)

    Returns:
        y (torch.Tensor): Shape (batch_size, dim, seq_len)
    """
    weights = normalized_dot_product(x, x)
    return weighted_avg(x, weights)
```

# 2. A self-attention module
Like you saw in the video lectures, self-attention is rarely used in the basic form we have created above.
Let's do the modifications needed to construct an actual transformer.

We will wrap it in a proper `torch.nn` module to create a building block that we can use in a network.
Creating your own module is actually not that common, frameworks like `pytorch` are built to be *modular* and we can often create very specific networks by combining standard modules. That is a good thing, since it enables us to express interesting models in a high-level interface and as a bonus, we build a model from well-tested and efficient parts.
With that said, you might find yourself in a situation (perhaps already in the project) where no off-the-shelf module suits your need and you have to create one yourself. View this latter part as an example/inspiration of how to construct a non-trivial custom module.

## Queries, keys and values
The self-attention is extended with three linear mappings $\Wq, \Wk, \Wv \in \R^{\dim \times \dim}$ .
These give us learnable parameters and make self-attention more flexible.
The three matrices map the input $\x_i$ into a query, key and value respectively:

\begin{align}
    \q_i = \Wq \x_i \\
    \k_i = \Wk \x_i \\
    \v_i = \Wv \x_i
\end{align}

First, we modify the self-attention by redefining the unnormalized weights (while reusing the notation):

\begin{align}
    \wu_{ji} = \q_j^T \k_i \Big{/} \sqrt{\dim}
\end{align}
The normalized weights are still obtained by applying the softmax function.

**[2 points]** Explain why we scale the dot product with the factor $1 / \sqrt{\dim}$.

**Your answer:** (fill in here)

Finally, the weighted average modified and is now based on the values $\v_j$, instead of on $\x_j$ directly:
$$
\y_i = \sum_{j} \v_j \w_{ji}.
$$

We can reuse our dot product calculation by simple *wrapping* it in a function that takes queries and keys as the argument:

```python
def query_key_weights(queries, keys):
    """Weights from query-key dot product.
    Softmax-normalised dot product weights
    Calculates weights for a batch of sequences of vectors.

    Args:
        queries (torch.Tensor): Shape (batch_size, dim, seq_len)
        keys (torch.Tensor): Shape (batch_size, dim, seq_len)

    Returns:
        weights (torch.Tensor): Shape (batch_size, seq_len, seq_len)
    """
    dim = queries.shape[2]
    queries = queries / (dim ** (1/4))
    keys    = keys / (dim ** (1/4))
    return normalized_dot_product(queries, keys)
```

## Multi-head self-attention

The model should be able to find different patterns in the input sequence, which is why we use multiple heads.

Now, we'll create the actual self-attention function, which includes multiple heads.
For implementation simplicity and efficiency we will do a version called *narrow* self-attention, where the input vector is split into parts and each attention head is applied to just one part of the vector.
Imagine that we have $\d = 64$ and four heads, then each head would operate on a vector with dimension $64 / 4 = 16$.

## Constructing the module
Below is an implementation of our self-attention module. We try to show you how a typical custom model looks like. Part of that is to do full vectorization (i.e. no loops). The result is a lot of manipulation of shapes and dimension order of intermediate tensors. It is not very readable and it is quite difficult to wrap your head around it but since you are likely to use and modify other peoples code (in the project or some later time), it is good that you get exposed to it now.

```python
class SelfAttention(torch_nn.Module):
    def __init__(self, dim, heads):
        """(Narrow) Self-attention module

        Args:
            dim (int): The full embedding dimension of the input vectors
            heads (int): The number of heads in the multi-head attention.
        """
        super().__init__()
        if not dim % heads == 0:
            raise ValueError(
                "The embedding dim. must be divisible by the number of heads for the vectorization to work."
            )
        self.dim = dim
        self.heads = heads
        part_dim = dim // heads
        # Linear maps for q, k and v
        self.Wq = torch_nn.Linear(part_dim, part_dim, bias=False)
        self.Wk = torch_nn.Linear(part_dim, part_dim, bias=False)
        self.Wv = torch_nn.Linear(part_dim, part_dim, bias=False)
        # Linear mapping to return to the original
        self.WO = torch_nn.Linear(heads * part_dim, dim)

    def forward(self, x):
        """Multi-headed self attention

        Each head operates on a part of the embedding, i.e. we have q, k and v with shape
        (batch_size, seq_length, heads, dim / heads)

        Args:
            x (Tensor): Input with shape (batch_size, seq_length, dim)
        """
        batch_size, seq_length, dim = x.shape
        part_dim = dim // self.heads
        x = x.reshape(batch_size, seq_length, self.heads, part_dim)

        keys = self.Wk(x)
        queries = self.Wq(x)
        values = self.Wv(x)

        keys = self._restructure_tensor(keys, batch_size, seq_length, part_dim)
        queries = self._restructure_tensor(queries, batch_size, seq_length, part_dim)
        values = self._restructure_tensor(values, batch_size, seq_length, part_dim)

        weights = query_key_weights(queries, keys)

        y_tilde = weighted_avg(values, weights)
        y_tilde = (
            y_tilde.transpose(2, 1)
            .reshape(batch_size, self.heads, seq_length, part_dim)
            .transpose(1, 2)
            .contiguous()
            .reshape(batch_size, seq_length, part_dim * self.heads)
        )
        return self.WO(y_tilde)

    def _restructure_tensor(self, x, batch_size, seq_length, part_dim):
        """ Reshaping q, k and v tensors

        For efficient vectorisation we stack the different heads in the batch_size dimension.
        Think of it as temporarily expanding the batch_size with every head.
        """
        return (
            x.transpose(1, 2)
            .contiguous()
            .reshape(batch_size * self.heads, seq_length, part_dim)
            .transpose(2, 1)
        )
```

# The transformer block

The majority of the implementation complexity is actually in the `SelfAttention` module. The transformer block is rather straight forward, it is just like the one described in the video lectures:

```python
class TransformerBlock(torch_nn.Module):
    """Transformer block"""

    def __init__(self, dim, heads):
        super().__init__()

        self.self_attention = SelfAttention(dim, heads)

        self.normalization_1 = torch_nn.LayerNorm(dim)
        self.normalization_2 = torch_nn.LayerNorm(dim)

        # The size of the hidden layer is a hyper-parameter,
        # but the consensus is that it should at least be larger than the input/output size
        self.feed_forward = torch_nn.Sequential(
            torch_nn.Linear(dim, 4 * dim),
            torch_nn.ReLU(),
            torch_nn.Linear(4 * dim, dim),
        )

    def forward(self, x):
        y = self.self_attention(x)
        # Note how the residual (skip) connections are implemented as simple addition.
        x = self.normalization_1(x + y)
        fed_forward = self.feed_forward(x)
        return self.normalization_2(fed_forward + x)
```

Now, we are done with the general module. To create an actual transformer yet we must choose an actual problem so that we can specify input, embedding and output.
Let's do that.

# 3. IMDB Classification

In this task we will consider classification of IMDB reviews. For pedagogical reasons, the classification will be based upon the functions you have built above. The idea is to clarify the connection between these functions and an application with real data. However, note that the standard approach to maximize performance is to use a well studied transformer architecture (e.g. BERT), download pre-trained weights and then fine-tune the parameters for the considered problem. Additional resources on how one would implement a transformer in practice are appended at the end of this notebook.

The purpose here is to build on the computer labs and to give you some inspiration for how to solve a general problem with `pytorch`. It will show you how to install additional python libraries (useful for the project) and some advice on how to construct a training/validation loop. We do not expect you to modify the code, **you don't even have to run it** if you feel that your cloud credits are starting to run low. However, you should read and understand the code; it will help you answer the questions at the end.

## The data

The IMDB data is provided by an external python module called `torchtext`.
You can add it to the dml conda environment with:
```
conda install -c pytorch torchtext
```
Make sure that you have activated the dml environment before your run it.

Processing text data can be tedious and error prone. For prototyping it is nice to use some third-party library which has done most of the work for you. You do not really need to focus on the data processing here, since it will be different for every task.

```python
from torchtext import data, datasets

TEXT = data.Field(lower=True, include_lengths=True, batch_first=True)
LABEL = data.Field(sequential=False)

def get_loaders(vocabulary_len, batch_size, device, split_ratio=0.8):
    """Load the IMDB data"""
    tdata, _ = datasets.IMDB.splits(TEXT, LABEL)
    train, test = tdata.split(split_ratio)

    TEXT.build_vocab(
        # We have to leave space for two special tokens.
        train, max_size=vocabulary_len - 2
    )
    LABEL.build_vocab(train)

    train_loader, test_loader = data.BucketIterator.splits(
        (train, test), batch_size=batch_size, device=device
    )
    return train_loader, test_loader

def view_example_text(index):
    """Helper function to look at a sample.

    The dataset is quite slow to load.
    """
    train, test = datasets.IMDB.splits(TEXT, LABEL)
    sample_text = train[index].text
    sample_label = train[index].label
    # Simply print the list of words, separated by space.
    print(" ".join(sample_text))
    print("\nLabel: ", sample_label)

view_example_text(118)
```

## The transformer
We will create a simple transformer that takes as input text in the form of a python list of words and which outputs a  probability vector over the two classes "pos" and "neg" (technically, the output will be the input to a log-softmax).
We make the simplest (and less memory efficient) version of position embedding as described in the video lectures.

```python
class Transformer(torch_nn.Module):
    def __init__(self, dim, heads, depth, seq_length, num_tokens, num_classes, device):
        super().__init__()
        self.device = device

        self.num_tokens = num_tokens

        self.pos_emb = torch_nn.Embedding(seq_length, dim)
        self.token_emb = torch_nn.Embedding(num_tokens, dim)

        transformer_blocks = []
        for _ in range(depth):
            transformer_blocks.append(TransformerBlock(dim=dim, heads=heads))

        # The Sequential wrapper is convenient when you want to repeat similar blocks.
        # A down-side is that it is harder access intermediate values for debugging.
        self.transformer_blocks = torch_nn.Sequential(*transformer_blocks)

        # The last part is problem specific. Here we want to map our transformer embeddings
        # to a probability distribution.
        # We will use a linear layer to produce log logits (the input to a  log-softmax function).
        self.output_map = torch_nn.Linear(dim, num_classes)

    def forward(self, x):
        """Transformer forward method

        Args:
            x Tensor(batch_size, seq_length): Word indices representing sequence of words.
        Returns:
            Tensor(batch_size, num_classes): Log logits
        """
        tokens = self.token_emb(x)
        batch_size, seq_length, dim = tokens.size()

        # Note that we create a completely new tensor which must be moved to the proper device.
        # This is why we must store the device in self.device.
        pos = torch.arange(seq_length, device=self.device)
        pos = self.pos_emb(pos)[None, :, :].expand(batch_size, seq_length, dim)

        x = tokens + pos
        x = self.transformer_blocks(x)

        x = self.output_map(x.mean(dim=1))
        return F.log_softmax(x, dim=1)
```

Now, for the train./val loop. This can be written in many ways but based on common misstakes in HA1, hints might be in order:

- Separate your code into smaller pieces, i.e. functions. It makes it easier to find bugs and easier to reuse code.
- Use separate functions to calculate metrics. If you want to calculate, say accuracy, during both training and validation, don't copy the code. Write one function and make sure that it works, then reuse it.
- Adding measurements to a running metrics can be tricky. Below is a solution that is a bit overkill but that is okay, since it is hard to use it incorrectly.

Note 1: the code below can be modified so that you can play around with it.

Note 2: timing this on Azure, a single epoch took ~5 min. Feel free to reduce the number of epochs or just study the code.

```python
from time import time

def train_epoch(model, train_loader, optimizer, scheduler, max_seq_len):
    """Train epoch"""
    train_loss = AccumulatingMetric()
    train_acc = AccumulatingMetric()
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_, label = batch.text[0], batch.label - 1

        input_ = _truncate_input(input_, max_seq_len)
        pred = model(input_)
        loss = F.nll_loss(pred, label)
        loss.backward()
        train_loss.add(loss.item())

        train_acc.add(accuracy(pred, label))

        # Gradient clipping is a way to ensure
        # torch_nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
    return train_loss.avg(), train_acc.avg()

def validate_epoch(model, val_loader, max_seq_len):
    val_loss = AccumulatingMetric()
    val_acc = AccumulatingMetric()
    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            input_, label = batch.text[0], batch.label - 1

            input_ = _truncate_input(input_, max_seq_len)
            pred = model(input_)
            val_loss.add(F.nll_loss(pred, label).item())

            val_acc.add(accuracy(pred, label))

    return val_loss.avg(), val_acc.avg()  # TODO: loss

def accuracy(pred, label):
    hard_pred = pred.argmax(1)
    return (hard_pred == label).float().mean().item()

def _truncate_input(input_, max_seq_len):
    if input_.size(1) > max_seq_len:
        input_ = input_[:, :max_seq_len]
    return input_

class AccumulatingMetric:
    """Accumulate samples of a metric and automatically keep track of the number of samples."""
    def __init__(self):
        self.metric = 0.0
        self.counter = 0

    def add(self, value):
        self.metric += value
        self.counter += 1

    def avg(self):
        return self.metric / self.counter

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_tokens = 50_000
max_length = 512
embedding_size = 128
num_heads = 8
num_classes = 2
depth = 6

model = Transformer(
    dim=embedding_size,
    heads=num_heads,
    depth=depth,
    seq_length=max_length,
    num_tokens=num_tokens,
    num_classes=num_classes,
    device=device)

model.to(device)

lr = 1e-4
lr_warmup = 1e4
num_epochs = 5
batch_size = 4

train_loader, test_loader = get_loaders(num_tokens, batch_size, device)

optimizer = torch.optim.Adam(lr=lr, params=model.parameters())
# A scheduler is a principled way of controlling (often decreasing) the learning rate as time progresses.
# Read more: https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lambda i: min(i / (lr_warmup / batch_size), 1.0)
)

print("Starting training")
for epoch in range(1, num_epochs + 1):
    start = time()
    train_loss, train_acc = train_epoch(model, train_loader, optimizer, scheduler, max_length)
    val_loss, val_acc = validate_epoch(model, test_loader, max_length)
    end = time()
    print(
        "Epoch: {}/{}: time: {:.1f}, train loss: {:.3f}, train acc: {:.3f}, val. loss {:.3f}, val. acc: {:.3f}".format(
            epoch, num_epochs, end - start, train_loss, train_acc, val_loss, val_acc
        )
    )
print("You have now trained a transformer!")

# I'm adding a ''# YOUR CODE HERE' tag so that the code above is not hidden when the assignment is generated.
# YOUR CODE HERE
```

## Theoretical questions

Now that you have gotten a practical feel for the transformer it is time to reflect on some of its important properties:

## Multi-Head Self Attention

**[3 points]** The above implementation of the transformer is using multi-head self-attention. Briefly describe which parts that are different in different heads in a multi-head self-attention layer.

**Your answer:** (fill in here)

**[3 points]** What are the benefits of using multi-head self-attention compared to a single-head self-attention block?

**Your answer:** (fill in here)

## Transformers and RNN

**[7 points]** Which are the main differences between a transformer and an RNN?
In particular, how do the differences make it easier to train a transformer, compared to an RNN?

**Your answer:** (fill in here)

**[7 points]** Self-attention maps sets to sets. It is an important part of what makes transformers so useful and general. Explain what this property means?

Ironically, this property is actually a bit of an issue when we want to process text (or any NLP problem).
Why, and how do we try to fix it?

**Your answer:** (fill in here)

# Wrapping up

The transformer architecture has become incredibly popular and has produced truly amazing results.
You should now have a good insight for how they can be implemented in `pytorch`. If you are interested, here are some more resources:

- https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html
- https://github.com/huggingface/transformers
- https://towardsdatascience.com/bert-text-classification-using-pytorch-723dfb8b6b5b


 # Prescreen results

```python
import matplotlib as mpl
import matplotlib.pyplot as plt
import os
import warnings
import numpy as np

from mpl_toolkits.axes_grid1 import make_axes_locatable

import nelpy as nel
import nelpy.plotting as npl

npl.setup()

get_ipython().run_line_magic('matplotlib', 'inline')

warnings.filterwarnings("ignore")

```

    /home/ckemere/Code/nelpy/nelpy/plotting/decoding.py:9: MatplotlibDeprecationWarning:
    The mpl_toolkits.axes_grid module was deprecated in Matplotlib 2.1 and will be removed two minor releases later. Use mpl_toolkits.axes_grid1 and mpl_toolkits.axisartist, which provide the same functionality instead.

```python
datadirs = ['/mnt/d/Data/Etienne/','/home/ckemere/Data2/DibaBuzsakiData/',
    'D:/Dropbox/neoReader/']

## conda install pandas=0.19.2

fileroot = next( (dir for dir in datadirs if os.path.isdir(dir)), None)

if fileroot is None:
    raise FileNotFoundError('datadir not found')

load_from_nel = True

# load from nel file:
if load_from_nel:
    jar = nel.load_pkl(fileroot + 'gor01vvp01pin01_processed_speed.nel')
    aux_data = jar.aux_data
    jar = nel.load_pkl(fileroot + 'gor01vvp01pin01_tables_speed.nel')
    df = jar.df
    df2 = jar.df2
    jar = nel.load_pkl(fileroot + 'score_all_sessions_5000_35000.nel') # needs to be modified to include pin01 bayes
    score_results = jar.results
    del jar

jar = nel.load_pkl('sessions-segments.nel')
sessions, segments = jar.sessions, jar.segments

```

```python
def plot_all_PBEs(bst, spiketrainarray, tuningcurve, tc_placecells, idx=None, title_str=None, size=None):
    if idx is not None:
        bst = bst[idx]
    st = spiketrainarray
    tc = tuningcurve
    tc_placecells = tc_placecells

    no = tc_placecells.get_peak_firing_order_ids()
    st.reorder_units_by_ids(no, inplace=True)

    st_cut = st[bst.support]
    st_cut._support = bst.support # hacky fix so that we can plot events out of order
    st_cut = nel.utils.collapse_time(st_cut)

    # decode neural activity
    posterior, bdries, mode_pth, mean_pth = nel.decoding.decode1D(bst=bst, ratemap=tc, xmax=310)

    if size == 'big':
        figwidth = 0.8 * bst.n_bins
    else:
        figwidth = 0.08 * bst.n_bins

    with npl.FigureManager(show=True, figsize=(figwidth,5)) as (fig, ax):
        npl.utils.skip_if_no_output(fig)

        pixel_width = 0.5

        npl.imagesc(x=np.arange(bst.n_bins), y=np.arange(311), data=posterior, cmap=plt.cm.Spectral_r, ax=ax)
        npl.utils.yticks_interval(310)
        npl.utils.no_yticks(ax)
        # plt.imshow(posterior, cmap=plt.cm.Spectral_r, interpolation='none', aspect='auto')
        ax.vlines(np.arange(bst.lengths.sum())-pixel_width, *ax.get_ylim(), lw=1, linestyle=':', color='0.8')
        ax.vlines(np.cumsum(bst.lengths)-pixel_width, *ax.get_ylim(), lw=1)

        ax.set_xlim(-pixel_width, bst.lengths.sum()-pixel_width)

        event_centers = np.insert(np.cumsum(bst.lengths),0,0)
        event_centers = event_centers[:-1] + bst.lengths/2 - 0.5

        ax.set_xticks(event_centers);
        if idx is not None:
            ax.set_xticklabels(idx);
        else:
            ax.set_xticklabels(np.arange(bst.n_epochs));

        npl.utils.no_xticks(ax)

        divider = make_axes_locatable(ax)
        if size == 'big':
            axRaster = divider.append_axes("top", size=1.5, pad=0)
        else:
            axRaster = divider.append_axes("top", size=1, pad=0)

        npl.rasterplot(st_cut, vertstack=True, ax=axRaster, lh=2)
        axRaster.set_xlim(st_cut.support.time.squeeze())
        bin_edges = np.linspace(st_cut.support.time[0,0],st_cut.support.time[0,1], bst.n_bins+1)
    #     axRaster.vlines(bin_edges, *ax.get_ylim(), lw=1, linestyle=':', color='0.2')
        axRaster.vlines(bin_edges[np.cumsum(bst.lengths)], *ax.get_ylim(), lw=1, color='0.2')
        npl.utils.no_xticks(axRaster)
        npl.utils.no_xticklabels(axRaster)
        npl.utils.no_yticklabels(axRaster)
        npl.utils.no_yticks(axRaster)
        ax.set_ylabel('position [cm]')
        ax.set_xlabel('time bins (20 ms)')
        if title_str:
            fig.suptitle(title_str)
        npl.utils.clear_left_right(axRaster)
        npl.utils.clear_top_bottom(axRaster)

```

```python
min_n_placecells = 20
min_n_PBEs = 30
df2_subset = df2[(df2.n_PBEs >= min_n_PBEs) & (df2.n_placecells >= min_n_placecells)]

```

```python
z_scores = np.array(df2_subset['prescreen_z'].values.tolist())
z_order = np.argsort(z_scores)[::-1] # ascending order

PlotAllSessions = False
if PlotAllSessions:
    for session, segment in zip(np.array(sessions)[z_order], np.array(segments)[z_order]):
        bst = aux_data[session][segment]['PBEs']
        st = aux_data[session][segment]['st_placecells']
        tc = aux_data[session][segment]['tc']
        tc_placecells = aux_data[session][segment]['tc_placecells']
        evt_order = np.argsort(score_results[session][segment]['scores_bayes_percentile']).tolist()[::-1]
        title_str = str(session) + ': ' + str(segment) + '; events ordered by decreasing Bayesian score percentile'
        plot_all_PBEs(bst, st, tc, tc_placecells, idx=evt_order, title_str=title_str)

```

 # Look at an example session

```python
session = '16-40-19'; segment = 'short'
# session = '1-22-43'; segment = 'short'
bst = aux_data[session][segment]['PBEs']
st = aux_data[session][segment]['st_placecells']
tc = aux_data[session][segment]['tc']
tc_placecells = aux_data[session][segment]['tc_placecells']
bst.n_epochs

scores_bayes_percentile = score_results[session][segment]['scores_bayes_percentile']
scores_hmm_percentile = score_results[session][segment]['scores_hmm_percentile']
```

```python
import warnings
warnings.filterwarnings('once')
```

### _Etienne has gone through and scored all the sequences in this session_

```python
#plot_all_PBEs(bst, st, tc, tc_placecells) # Plot all the sequences in the session
```

```python
goodseqs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
              21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
              40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58,
              59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77,
              78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96,
              97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112,
              113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127,
              128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142,
              143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157,
              158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172,
              173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,
              188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202,
              203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217,
              218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232,
              233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247,
              248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262,
              263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276]

badseqs = [0,2,3,14,34,47,55,82,88,108,126,135,137,153,180,181,184,185,207,229,237,251,
          256,259,261,263,272,72,86,91,169,172,204,208,220,241,240,12,81,83,196]

uncertainseqs = [28,27,44,51,61,63,80,92,93,98,66,136,156,158,170,216,233,246,230,56,46,
                 45,37,48,79,99,110,114,116,120,125,130,138,161,163,164,183,197]

statseqs = [1,5,7,9,10,13,16,21,22,31,33,39,57,58,59,76,77,97,104,118,123,140,142,148,151,
           166,167,168,179,194,200,201,212,223,224,225,245,242,243,239,273,274,275,276,73,
           84,100,124,254,139,189,248,268,232,252,244,236,64,67,68,65,43,42,4,29,85,87,102,
            117,119,134,186,209]

goodseqs = list(set(goodseqs) - set(statseqs) - set(badseqs) - set(uncertainseqs))
uncertainseqs = list(set(uncertainseqs) - set(badseqs))
statseqs = list(set(statseqs) - set(badseqs) - set(uncertainseqs))
```

### Let's look at an example from the good sequences
The plot_all_PBEs function plots the spike rasters ordered by place cell locations above the Bayesian decoded linear position for a list of sequences. Lets pick a set of sequences that look "good" (replay is of a nice linear path along the track), but have scores in the 98-99 percentile range rather than the max (100). Presumably these sequences will be marginally more likely to be affected by changes.

```python
niceseq = [54, 74, 90, 197, 203, 213, 227]
print(scores_bayes_percentile[niceseq])
```

    [100.    99.34  99.98  99.16  99.22  98.86  98.34]

```python
plot_all_PBEs(bst, st, tc, tc_placecells, idx=niceseq, size='big')

```

#### How do we simulate the effect of SWR firing patterns being different in different conditions?

One simple test is to take the binned spikes and, for example, triple the number of spikes. This sharpens the posteriors as shown below.

```python
bst2 = bst * 3 # simulate a doubling of firing rates during ripples
plot_all_PBEs(bst2, st, tc, tc_placecells, idx=niceseq, size='big')

```

How we choose the tuning curves for replay analysis is already a bit suspect. Usually, for convenience, we just say that the firing rate is scaled by a factor that is equivalent to the ratio of the bin sizes from running to ripple periods. But when we consider in detail the question of whether changes in the pattern of spiking during ripples might lead to differences in replay detection, we should consider this more carefully. For some position $x$ and the position-specific firing rate map $\boldsymbol{\lambda}_x$, if we observe a spiking vector $\mathbf{k}$, the posterior likelihood is given by:

$$ LL(x) = \mathbf{k}^T \log(\boldsymbol{\lambda}_x) - \mathbf{1}^T \mathbf{\lambda}_x + C $$

where $C$ is a constant that is independent of $x$. If we scale the rate by a factor $\alpha$, the new likelihood is:

$$\begin{align}
LL(x) &= \mathbf{k}^T \log(\alpha \boldsymbol{\lambda}_x) - \alpha (\mathbf{1}^T \mathbf{\lambda}_x) + C \\
      &= \mathbf{k}^T \log(\boldsymbol{\lambda}_x) - \alpha (\mathbf{1}^T \mathbf{\lambda}_x) + C'
\end{align}$$

where $C'$ has absorbed the position-independent term $\mathbf{k}^T \log(\alpha)$. What's interesting about this is that the effect of tuning curve scaling is on the "prior" portion of the likelihood, the part that is independent of firing rates. This relates to the bigger picture idea that experimental conditions might bias results in unexpected ways - here, the bias is that place fields will not uniformly tile the space, so the prior may have some effect on replay score.

So below, let's plot what happens if we take the tuning curves and scale all their firing rates by a factor of three. You'll see it looks very similar.

```python
tc2 = tc * 0.333 # simulate a different firing rate map
plot_all_PBEs(bst, st, tc2, tc_placecells, idx=niceseq, size='big')

ratemap = tc.reorder_units_by_ids(bst.unit_ids)
prior = -ratemap.ratemap.sum(axis=0)*bst.ds
plt.plot(ratemap.bin_centers, prior)
plt.gca().set_title('Map of Prior Probability of Track');
```

#### What is the effect on replay scores?

We'll use the Davidson-style replay detection to compare scores between different scenarios.

```python
print('Running replay analsysis on {} SWR sequences'.format(bst.n_intervals))
```

    Running replay analsysis on 277 SWR sequences

```python
from nelpy import analysis

scores1, shuffled_scores1, percentiles1 = \
    nel.analysis.replay.score_Davidson_final_bst_fast(bst=bst,
        tuningcurve=tc,w=3,n_shuffles=1000, n_samples=10000)
```

    /home/ckemere/anaconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject

```python
scores2, shuffled_scores2, percentiles2 = \
    nel.analysis.replay.score_Davidson_final_bst_fast(bst=bst,
        tuningcurve=tc2,w=3,n_shuffles=1000, n_samples=10000)
```

```python
scores3, shuffled_scores3, percentiles3 = \
    nel.analysis.replay.score_Davidson_final_bst_fast(bst=bst2,
        tuningcurve=tc,w=3,n_shuffles=1000, n_samples=10000)

```

```python
import seaborn as sns

# Plot score distributions for each event

if False:
    to_plot = [(scores1, shuffled_scores1), (scores2, shuffled_scores2), (scores3, shuffled_scores3)]

    fig, axs = plt.subplots(nrows=len(to_plot), ncols=len(scores1), figsize=(4 * len(scores1), 8))
    plt.subplots_adjust(hspace=0.5)

    npl.setup(style='whitegrid')
    for idx, (sc, sh) in enumerate(to_plot):
        for i in range(len(scores1)):
            ax = axs[idx,i]
            sns.distplot(sh[i,:],ax=ax)
            ax.stem([sc[i]],[1],markerfmt='C1o')
            if (idx > 0):
                ax.set_xlim(axs[0,i].get_xlim())
    axs[0,1].set_title('Original');
    axs[1,1].set_title('0.1 x TC');
    axs[2,1].set_title('2 x Spikes');

    print(np.mean(percentiles1-percentiles2), np.std(percentiles1-percentiles2))
    print(np.mean(percentiles1-percentiles3), np.std(percentiles1-percentiles3))
```

```python
print('Mean percentile score for events:\n')
print('   Original     TC/3        FR*3')
print(np.vstack([percentiles1,percentiles2,percentiles3]).mean(axis=1))
```

    Mean percentile score for events:

       Original     TC/3        FR*3
    [97.74963899 98.28501805 97.41119134]

```python
sns.boxplot(data=np.vstack([percentiles1,percentiles2,percentiles3]).T)
```

```python
from scipy.stats import wilcoxon

print('Wilcoxon Tests')

statistic, pval = wilcoxon(percentiles1, percentiles2)
print('Original vs TC/3: ', statistic, pval)

statistic, pval = wilcoxon(percentiles1, percentiles3)
print('Original vs FR*3: ', statistic, pval)
```

    Wilcoxon Tests
    Original vs TC/3:  4324.0 0.011929552972645814
    Original vs FR*3:  4807.5 0.04839551607686778

```python
print('Fraction of events significant at 99%:\n')
print('  Original     TC/3     FR*3')
print((np.vstack([percentiles1,percentiles2,percentiles3]) > 99).mean(axis=1))
```

    Fraction of events significant at 99%:

      Original     TC/3     FR*3
    [0.74729242 0.77256318 0.68592058]

## Conclusion

When we look at all the sequences, it looks like both the tuning curve scaling and the firing rate change have small but potentially significant effects on the replay scores - lowering the tuning curve firing rates by a factor of 1/3rd slightly increases the mean score percentile and increasing the firing rates by a factor of 3 during ripples slightly decreases the average. This points to the broader picture that changes in firing patterns during ripples which are not accompanied by changes in the tuning curves might lead to changes in the statistics of fraction of "significant" replay or other aggregations of replay scores. We made fairly large changes to the model or data, but these have fairly small effects on "the number of significant events".

Replay Scores and Tuning Curve Scaling.ipynb