"In this TD, you must modify this notebook to complete the code (**# TO DO comments**) and complete the **proposed experiments**. To do this,\n",
"\n",
"1. Fork this repository\n",
"2. Clone your forked repository on your local computer\n",
"3. Add your code and answer the questions\n",
"4. Commit and push regularly\n",
"\n",
"**The last commit is due on Sunday, 14th January 2024**. Later commits will not be taken into account.\n",
"\n",
"As the computation is heavy, particularly during training, we encourage you to use a GPU. If your laptob is not equiped, you may use one of these remote jupyter servers, where you can select the execution on GPU :\n",
"This server is accessible within the campus network. If outside, you need to use a VPN. Before executing the notebook, select the kernel \"Python PyTorch\" to run it on GPU and have access to PyTorch module.\n",
"Before executing the notebook, select the execution on GPU : \"Exécution\" Menu -> \"Modifier le type d'exécution\" and select \"T4 GPU\"."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6oHNXwKGc4wh"
},
"source": [
"### Goal of the TD\n",
"\n",
"Transformers have been introduced by [Vaswani et al. in 2017](https://arxiv.org/abs/1706.03762) in the context of NLP (Natural Language Processing), and particulary for Machine Translation.\n",
"\n",
"Its great success has led to its adaptation to various applications, including image classification. In this trend, [Dosovitskiy et al. in 2020](https://arxiv.org/abs/2010.11929) have proposed Vision Transformers (ViT) that we will study and implement from scratch in this TD.\n",
"\n",
"The principle is illustrated in the following picture from this paper.\n",
"First, an input image is “cut” into sub-images equally sized.\n",
"\n",
"Each such sub-image goes through a linear embedding. From then, each sub-image becomes a one-dimensional vector.\n",
"\n",
"A positional embedding is then added to these vectors (tokens). The positional embedding allows the network to know where each sub-image is positioned originally in the image. Without this information, the network would not be able to know where each such image would be placed, leading to potentially wrong predictions.\n",
"\n",
"These tokens are then passed, together with a special classification token, to the transformer encoders blocks, were each is composed of : A Layer Normalization (LN), followed by a Multi-head Self Attention (MSA) and a residual connection. Then a second LN, a Multi-Layer Perceptron (MLP), and again a residual connection. These blocks are connected back-to-back.\n",
"\n",
"Finally, a classification MLP head is used for the final classification only on the special classification token, which by the end of this process has global information about the picture.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0NEnfO0Bc4wi"
},
"source": [
"### Implementation of the ViT model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kqJ6S1EFc4wj"
},
"source": [
"First, we import the required modules."
]
},
{
"cell_type": "code",
"execution_count": 115,
"metadata": {
"id": "wEmbaOA4Okuo"
},
"outputs": [],
"source": [
"# Import modules\n",
"import numpy as np\n",
"import torch\n",
"import torch.nn as nn\n",
"from torch.nn import CrossEntropyLoss\n",
"from torch.optim import Adam\n",
"from torch.utils.data import DataLoader\n",
"from torchvision.datasets.mnist import MNIST\n",
"from torchvision.transforms import ToTensor"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eOXWp-N2c4wl"
},
"source": [
"For this first experiment, we will use the MNIST dataset that contains 28x28 binary pixels images of hand-written digits ([0–9])."
"The transformer encoder was originally developed with sequence data in mind, such as English sentences. However, as an image is not a sequence, we need to “sequencify” an image. To do this, we break it into multiple sub-images and map each sub-image to a vector.\n",
"\n",
"We do so by simply reshaping our input, which has size (N, C, H, W), where N is the batch size, C the number of channels and (H,W) the image dimension. In the case of MNIST, dimensions are (N, 1, 28, 28). The target dimension is (N, #Patches, Patch dimensionality), where the dimensionality of a patch is adjusted accordingly.\n",
"\n",
"In this example, we break each (1, 28, 28) into 7x7 patches (hence, each of size 4x4). That is, we are going to obtain 7x7=49 sub-images out of a single image.\n",
"\n",
"Thus, we reshape input (N, 1, 28, 28) to (N, PxP, C x H/P x W/P) = (N, 49, 16)\n",
"\n",
"Notice that, while each patch is a picture of size 1x4x4, we flatten it to a 16-dimensional vector. Also, in this case, we only had a single color channel. If we had multiple color channels, those would also have been flattened into the vector."
]
},
{
"cell_type": "code",
"execution_count": 118,
"metadata": {
"id": "fxhHKKDFOoHp"
},
"outputs": [],
"source": [
"def patchify(images, n_patches):\n",
" n, c, h, w = images.shape\n",
"\n",
" assert h == w, \"Patchify method is implemented for square images only\"\n",
"\n",
" patches = torch.zeros(n, n_patches**2, h * w * c // n_patches**2)\n",
" patch_size = h // n_patches\n",
"\n",
" for idx, image in enumerate(images):\n",
" for i in range(n_patches):\n",
" for j in range(n_patches):\n",
" patch = image[\n",
" :,\n",
" i * patch_size : (i + 1) * patch_size,\n",
" j * patch_size : (j + 1) * patch_size,\n",
" ]\n",
" patches[idx, i * n_patches + j] = patch.flatten()\n",
"Now that we have our flattened patches, we can map each of them through a Linear mapping. While each patch was a 4x4=16 dimensional vector, the linear mapping can map to any arbitrary vector size. Thus, we will use for this a parameter `hidden_d` for \"hidden dimension\".\n",
"\n",
"In this example, we will use a hidden dimension of 8, but in principle, any number can be put here. We will thus be mapping each 16-dimensional patch to an 8-dimensional patch.\n"
"Positional encoding allows the model to understand where each patch would be placed in the original image. While it is theoretically possible to learn such positional embeddings, previous work by [Vaswani et al. in 2017](https://arxiv.org/abs/1706.03762) suggests that we can just add sines and cosines waves.\n",
"\n",
"In particular, positional encoding adds high-frequency values to the first dimensions and lower-frequency values to the latter dimensions.\n",
"\n",
"In each sequence, for token i we add to its j-th coordinate the following value:\n",
"\n",
"$$p_{i,j}=\\left\\{\\begin{matrix}\n",
"sin(\\frac{i}{10000^{\\frac{j-1}{d_{emb-dim}}}}) \\:if \\:j \\: is \\: even\\\\\n",
"cos(\\frac{i}{10000^{\\frac{j-1}{d_{emb-dim}}}})\\:if\\: j \\: is \\: odd\n",
"This positional embedding is a function of the number of elements in the sequence and the dimensionality of each element. Thus, it is always a 2-dimensional tensor or “rectangle”.\n",
"\n",
"Here is a simple function that, given the number of tokens and the dimensionality of each of them, outputs a matrix where each coordinate (i,j) is the value to be added to token i in dimension j.\n",
"\n",
"This positional encoding is added to our model after the linear mapping and the addition of the class token."
"The objective is now that, for a single image, each patch has to be updated based on some similarity measure with the other patches. We do so by linearly mapping each patch (that is now an 8-dimensional vector in our example) to 3 distinct vectors: q, k, and v (query, key, value).\n",
"\n",
"Then, for a single patch, we are going to compute the dot product between its q vector with all of the k vectors, divide by the square root of the dimensionality of these vectors (sqrt(8)), softmax these so-called attention cues, and finally multiply each attention cue with the v vectors associated with the different k vectors and sum all up.\n",
"\n",
"In this way, each patch assumes a new value that is based on its similarity (after the linear mapping to q, k, and v) with other patches. This whole procedure, however, is carried out H times on H sub-vectors of our current 8-dimensional patches, where H is the number of Heads.\n",
"\n",
"Once all results are obtained, they are concatenated together. Finally, the result is passed through a linear layer (for good measure).\n",
"\n",
"The intuitive idea behind attention is that it allows modeling the relationship between the inputs. What makes a ‘0’ a zero are not the individual pixel values, but how they relate to each other.\n",
"\n",
"This is implemented in the MSA class:"
]
},
{
"cell_type": "code",
"execution_count": 125,
"metadata": {
"id": "CIoyR-QsOruC"
},
"outputs": [],
"source": [
"class MSA(nn.Module):\n",
" def __init__(self, d, n_heads=2):\n",
" super().__init__()\n",
" self.d = d\n",
" self.n_heads = n_heads\n",
"\n",
" assert d % n_heads == 0, f\"Can't divide dimension {d} into {n_heads} heads\"\n",
"\n",
" d_head = int(d / n_heads)\n",
" self.q_mappings = nn.ModuleList(\n",
" [nn.Linear(d_head, d_head) for _ in range(self.n_heads)]\n",
" )\n",
" self.k_mappings = nn.ModuleList(\n",
" [nn.Linear(d_head, d_head) for _ in range(self.n_heads)]\n",
" )\n",
" self.v_mappings = nn.ModuleList(\n",
" [nn.Linear(d_head, d_head) for _ in range(self.n_heads)]\n",
" )\n",
" self.d_head = d_head\n",
" self.softmax = nn.Softmax(dim=-1)\n",
"\n",
" def forward(self, sequences):\n",
" # Sequences has shape (N, seq_length, token_dim)\n",
" # We go into shape (N, seq_length, n_heads, token_dim / n_heads)\n",
" # And come back to (N, seq_length, item_dim) (through concatenation)\n",
" return torch.cat([torch.unsqueeze(r, dim=0) for r in result])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vomINDMpc4wq"
},
"source": [
"Notice that, for each head, we create distinct Q, K, and V mapping functions (square matrices of size 4x4 in our example).\n",
"\n",
"$$\\left\\{\\begin{matrix}\n",
"Q=W^Q.I\n",
"\\\\\n",
"K=W^K.I\n",
"\\\\\n",
"V=W^V.I\n",
"\\\\\n",
"A=K^T.Q\n",
"\\\\\n",
"\\hat{A}=softmax(A)\n",
"\\\\\n",
"O=V.\\hat{A}\n",
"\\end{matrix}\\right.$$\n",
"\n",
"Since our inputs will be sequences of size (N, 50, 8), and we only use 2 heads, we will at some point have an (N, 50, 2, 4) tensor, use a nn.Linear(4, 4) module on it, and then come back, after concatenation, to an (N, 50, 8) tensor.\n",
"\n",
"Also notice that using loops is not the most efficient way to compute the multi-head self-attention, but it makes the code much clearer for learning."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gG4UKx-Ac4wr"
},
"source": [
"### Transformer Encoder Blocks\n",
"\n",
"The next step is to create the transformer encoder block class.\n",
"\n",
"Layer normalization (LN) is a popular block that, given an input, subtracts its mean and divides by the standard deviation. It is applied to the last dimension only. We can thus make each of our 50x8 matrices (representing a single sequence) have mean 0 and std 1. After we run our (N, 50, 8) tensor through LN, we still get the same dimensionality.\n",
"\n",
"Also, We will be using residual connection that consists in adding the original input to the result of some computation. This, intuitively, allows a network to become more powerful while also preserving the set of possible functions that the model can approximate.\n",
"\n",
"We will add a residual connection that will add our original (N, 50, 8) tensor to the (N, 50, 8) obtained after LN and MSA.\n",
"\n",
"Next is to add a simple residual connection between what we already have and what we get after passing the current tensor through another LN and an MLP. The MLP is composed of two layers, where the hidden layer typically is four times as big (this is a parameter).\n",
"\n",
"The transformer encoder block class (which will be a component of the future ViT class) is thus as follows:"
"Now that the encoder block is ready, we just need to insert it in our bigger ViT model which is responsible for patchifying before the transformer blocks, and carrying out the classification after.\n",
"\n",
"To help classification, we will use an additional **classification token** to the input sequence. This is a special token that we add to our model that has the role of capturing information about the other tokens. This will happen with the MSA block. When information about all other tokens will be present here, we will be able to classify the image using only this special token. The initial value of the special token (the one fed to the transformer encoder) is a parameter of the model that needs to be learned.\n",
"\n",
"Thus, we will add a parameter to our model and convert our (N, 49, 8) tokens tensor to an (N, 50, 8) tensor (we add the special token to each sequence).\n",
"\n",
"We could have an arbitrary number of transformer blocks. In this example, to keep it simple, I will use only 2. We also add a parameter to know how many heads does each encoder block will use.\n",
"\n",
"Finally, we can extract just the classification token (first token) out of our N sequences, and use each token to get N classifications.\n",
"\n",
"Since we decided that each token is an 8-dimensional vector, and since we have 10 possible digits, we can implement the classification MLP as a simple 8x10 matrix, activated with the SoftMax function.\n",
"\n",
"The output of our model shoud be an (N, 10) tensor."
"1. Adapt the code to apply the ViT model on CIFAR dataset.\n",
"2. Make use of a validation set to evaluate overfitting.\n",
"3. Evaluate the model with a dimension of 16 for the tokens and 4 encoder blocks."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python PyTorch 1.7.0",
"language": "python",
"name": "pytorch-1.7.0"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.5"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
\ No newline at end of file
%% Cell type:markdown id: tags:
# TD3: Vision Transformer (ViT)
In this TD, you must modify this notebook to complete the code (**# TO DO comments**) and complete the **proposed experiments**. To do this,
1. Fork this repository
2. Clone your forked repository on your local computer
3. Add your code and answer the questions
4. Commit and push regularly
**The last commit is due on Sunday, 14th January 2024**. Later commits will not be taken into account.
As the computation is heavy, particularly during training, we encourage you to use a GPU. If your laptob is not equiped, you may use one of these remote jupyter servers, where you can select the execution on GPU :
This server is accessible within the campus network. If outside, you need to use a VPN. Before executing the notebook, select the kernel "Python PyTorch" to run it on GPU and have access to PyTorch module.
Before executing the notebook, select the execution on GPU : "Exécution" Menu -> "Modifier le type d'exécution" and select "T4 GPU".
%% Cell type:markdown id: tags:
### Goal of the TD
Transformers have been introduced by [Vaswani et al. in 2017](https://arxiv.org/abs/1706.03762) in the context of NLP (Natural Language Processing), and particulary for Machine Translation.
Its great success has led to its adaptation to various applications, including image classification. In this trend, [Dosovitskiy et al. in 2020](https://arxiv.org/abs/2010.11929) have proposed Vision Transformers (ViT) that we will study and implement from scratch in this TD.
The principle is illustrated in the following picture from this paper.
First, an input image is “cut” into sub-images equally sized.
Each such sub-image goes through a linear embedding. From then, each sub-image becomes a one-dimensional vector.
A positional embedding is then added to these vectors (tokens). The positional embedding allows the network to know where each sub-image is positioned originally in the image. Without this information, the network would not be able to know where each such image would be placed, leading to potentially wrong predictions.
These tokens are then passed, together with a special classification token, to the transformer encoders blocks, were each is composed of : A Layer Normalization (LN), followed by a Multi-head Self Attention (MSA) and a residual connection. Then a second LN, a Multi-Layer Perceptron (MLP), and again a residual connection. These blocks are connected back-to-back.
Finally, a classification MLP head is used for the final classification only on the special classification token, which by the end of this process has global information about the picture.
%% Cell type:markdown id: tags:
### Implementation of the ViT model
%% Cell type:markdown id: tags:
First, we import the required modules.
%% Cell type:code id: tags:
``` python
# Import modules
importnumpyasnp
importtorch
importtorch.nnasnn
fromtorch.nnimportCrossEntropyLoss
fromtorch.optimimportAdam
fromtorch.utils.dataimportDataLoader
fromtorchvision.datasets.mnistimportMNIST
fromtorchvision.transformsimportToTensor
```
%% Cell type:markdown id: tags:
For this first experiment, we will use the MNIST dataset that contains 28x28 binary pixels images of hand-written digits ([0–9]).
The transformer encoder was originally developed with sequence data in mind, such as English sentences. However, as an image is not a sequence, we need to “sequencify” an image. To do this, we break it into multiple sub-images and map each sub-image to a vector.
We do so by simply reshaping our input, which has size (N, C, H, W), where N is the batch size, C the number of channels and (H,W) the image dimension. In the case of MNIST, dimensions are (N, 1, 28, 28). The target dimension is (N, #Patches, Patch dimensionality), where the dimensionality of a patch is adjusted accordingly.
In this example, we break each (1, 28, 28) into 7x7 patches (hence, each of size 4x4). That is, we are going to obtain 7x7=49 sub-images out of a single image.
Thus, we reshape input (N, 1, 28, 28) to (N, PxP, C x H/P x W/P) = (N, 49, 16)
Notice that, while each patch is a picture of size 1x4x4, we flatten it to a 16-dimensional vector. Also, in this case, we only had a single color channel. If we had multiple color channels, those would also have been flattened into the vector.
%% Cell type:code id: tags:
``` python
defpatchify(images,n_patches):
n,c,h,w=images.shape
asserth==w,"Patchify method is implemented for square images only"
ax.imshow(patch[0],cmap="gray")# Utiliser la composante 0 pour niveaux de gris
ax.axis("off")
plt.show()
n_patches=4
image_plat=image.unsqueeze(0)
patches=patchify(image_plat,n_patches)
# Affichez les patches
display_patches(patches[0],n_patches,(1,7,7))
```
%% Output
%% Cell type:markdown id: tags:
### Linear embedding
Now that we have our flattened patches, we can map each of them through a Linear mapping. While each patch was a 4x4=16 dimensional vector, the linear mapping can map to any arbitrary vector size. Thus, we will use for this a parameter `hidden_d` for "hidden dimension".
In this example, we will use a hidden dimension of 8, but in principle, any number can be put here. We will thus be mapping each 16-dimensional patch to an 8-dimensional patch.
Positional encoding allows the model to understand where each patch would be placed in the original image. While it is theoretically possible to learn such positional embeddings, previous work by [Vaswani et al. in 2017](https://arxiv.org/abs/1706.03762) suggests that we can just add sines and cosines waves.
In particular, positional encoding adds high-frequency values to the first dimensions and lower-frequency values to the latter dimensions.
In each sequence, for token i we add to its j-th coordinate the following value:
$$p_{i,j}=\left\{\begin{matrix}
sin(\frac{i}{10000^{\frac{j-1}{d_{emb-dim}}}}) \:if \:j \: is \: even\\
cos(\frac{i}{10000^{\frac{j-1}{d_{emb-dim}}}})\:if\: j \: is \: odd
This positional embedding is a function of the number of elements in the sequence and the dimensionality of each element. Thus, it is always a 2-dimensional tensor or “rectangle”.
Here is a simple function that, given the number of tokens and the dimensionality of each of them, outputs a matrix where each coordinate (i,j) is the value to be added to token i in dimension j.
This positional encoding is added to our model after the linear mapping and the addition of the class token.
The objective is now that, for a single image, each patch has to be updated based on some similarity measure with the other patches. We do so by linearly mapping each patch (that is now an 8-dimensional vector in our example) to 3 distinct vectors: q, k, and v (query, key, value).
Then, for a single patch, we are going to compute the dot product between its q vector with all of the k vectors, divide by the square root of the dimensionality of these vectors (sqrt(8)), softmax these so-called attention cues, and finally multiply each attention cue with the v vectors associated with the different k vectors and sum all up.
In this way, each patch assumes a new value that is based on its similarity (after the linear mapping to q, k, and v) with other patches. This whole procedure, however, is carried out H times on H sub-vectors of our current 8-dimensional patches, where H is the number of Heads.
Once all results are obtained, they are concatenated together. Finally, the result is passed through a linear layer (for good measure).
The intuitive idea behind attention is that it allows modeling the relationship between the inputs. What makes a ‘0’ a zero are not the individual pixel values, but how they relate to each other.
This is implemented in the MSA class:
%% Cell type:code id: tags:
``` python
classMSA(nn.Module):
def__init__(self,d,n_heads=2):
super().__init__()
self.d=d
self.n_heads=n_heads
assertd%n_heads==0,f"Can't divide dimension {d} into {n_heads} heads"
Notice that, for each head, we create distinct Q, K, and V mapping functions (square matrices of size 4x4 in our example).
$$\left\{\begin{matrix}
Q=W^Q.I
\\
K=W^K.I
\\
V=W^V.I
\\
A=K^T.Q
\\
\hat{A}=softmax(A)
\\
O=V.\hat{A}
\end{matrix}\right.$$
Since our inputs will be sequences of size (N, 50, 8), and we only use 2 heads, we will at some point have an (N, 50, 2, 4) tensor, use a nn.Linear(4, 4) module on it, and then come back, after concatenation, to an (N, 50, 8) tensor.
Also notice that using loops is not the most efficient way to compute the multi-head self-attention, but it makes the code much clearer for learning.
%% Cell type:markdown id: tags:
### Transformer Encoder Blocks
The next step is to create the transformer encoder block class.
Layer normalization (LN) is a popular block that, given an input, subtracts its mean and divides by the standard deviation. It is applied to the last dimension only. We can thus make each of our 50x8 matrices (representing a single sequence) have mean 0 and std 1. After we run our (N, 50, 8) tensor through LN, we still get the same dimensionality.
Also, We will be using residual connection that consists in adding the original input to the result of some computation. This, intuitively, allows a network to become more powerful while also preserving the set of possible functions that the model can approximate.
We will add a residual connection that will add our original (N, 50, 8) tensor to the (N, 50, 8) obtained after LN and MSA.
Next is to add a simple residual connection between what we already have and what we get after passing the current tensor through another LN and an MLP. The MLP is composed of two layers, where the hidden layer typically is four times as big (this is a parameter).
The transformer encoder block class (which will be a component of the future ViT class) is thus as follows:
%% Cell type:code id: tags:
``` python
classViTBlock(nn.Module):
def__init__(self,hidden_d,n_heads,mlp_ratio=4):
super().__init__()
self.hidden_d=hidden_d
self.n_heads=n_heads
self.norm1=nn.LayerNorm(hidden_d)
self.mhsa=MSA(hidden_d,n_heads)
self.norm2=nn.LayerNorm(hidden_d)
self.mlp=nn.Sequential(
nn.Linear(hidden_d,mlp_ratio*hidden_d),
nn.GELU(),
nn.Linear(mlp_ratio*hidden_d,hidden_d),
)
defforward(self,x):
m1=x
x=self.norm1(x)
x=m1+self.mhsa(x)
m2=x
x=self.norm2(x)
x=m2+self.mlp(x)
returnx
```
%% Cell type:markdown id: tags:
### ViT model
Now that the encoder block is ready, we just need to insert it in our bigger ViT model which is responsible for patchifying before the transformer blocks, and carrying out the classification after.
To help classification, we will use an additional **classification token** to the input sequence. This is a special token that we add to our model that has the role of capturing information about the other tokens. This will happen with the MSA block. When information about all other tokens will be present here, we will be able to classify the image using only this special token. The initial value of the special token (the one fed to the transformer encoder) is a parameter of the model that needs to be learned.
Thus, we will add a parameter to our model and convert our (N, 49, 8) tokens tensor to an (N, 50, 8) tensor (we add the special token to each sequence).
We could have an arbitrary number of transformer blocks. In this example, to keep it simple, I will use only 2. We also add a parameter to know how many heads does each encoder block will use.
Finally, we can extract just the classification token (first token) out of our N sequences, and use each token to get N classifications.
Since we decided that each token is an 8-dimensional vector, and since we have 10 possible digits, we can implement the classification MLP as a simple 8x10 matrix, activated with the SoftMax function.
The output of our model shoud be an (N, 10) tensor.
