├── 0_intro └── README.md ├── 1_vision ├── README.md ├── vision_tutorial.ipynb └── vision_tutorial_solutions.ipynb ├── 2_mechanistic_interpretability ├── README.md └── mechanistic_interpretability_tutorial.ipynb ├── 3_rl ├── README.md ├── reinforcement_learning_tutorial.ipynb └── reinforcement_learning_tutorial_solutions.ipynb ├── 4_flow_matching_drug_discovery ├── README.md ├── drug_discovery.ipynb ├── drug_discovery_solutions.ipynb ├── flow_matching.ipynb └── flow_matching_solutions.ipynb └── README.md /0_intro/README.md: -------------------------------------------------------------------------------- 1 | # [[EEML2025](https://www.eeml.eu)] Tutorial 0: Introduction to Colab and PyTorch 2 | 3 | **Authors:** Mandana Samiei and Teodor Szente 4 | 5 | --- 6 | 7 | In this tutorial, we will learn how to use Colab notebook and train a simple model in PyTorch. 8 | 9 | ### Outline 10 | 11 | - [Tutorial 1] Intro to colab 12 | - [Tutorial 2] Intro to PyTorch (MLP) 13 | - [Tutorial 3] Advanced PyTorch (CNN) 14 | 15 | 16 | ### Notebooks 17 | 18 | Tutorial 1: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ocFPgLnE1X7YypI35x835AVhAj-gZulB?usp=sharing) 19 | 20 | Tutorial 2: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1B8JQGsoXJcZQ9n4fLivyOQTo27LY1ALd?usp=sharing) 21 | 22 | Tutorial 3: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ekLHmSVpuHaFHscZL_CrvnJgNaQpKTbi?usp=sharing) 23 | 24 | --- 25 | -------------------------------------------------------------------------------- /1_vision/README.md: -------------------------------------------------------------------------------- 1 | # [[EEML2025](https://www.eeml.eu)] Tutorial 1: Training Vision Transformer (ViT) encoders from scratch 2 | 3 | **Authors:** Liliane Momeni and Nikhil Parthasarathy 4 | 5 | 6 | --- 7 | 8 | This tutorial is designed to provide a hands-on understanding of how to train powerful Vision Transformer (ViT) encoders from scratch 🚀. Part 1 will recap why ViT encoders are important, before diving into the three fundamental training paradigms. Part 2 will cover the classic fully supervised approach using labeled data. Part 3 will explore the self-supervised paradigm through Masked Autoencoders, where models learn rich features from the data itself without labels. Finally, Part 4 covers contrastive learning with noisy supervision, using CLIP as our case study. 9 | 10 | ### Outline: 11 | 12 | - Part 1: Introduction to Vision Transformer encoders 13 | - Part 2: Fully supervised ViT for Image Classification 14 | - Part 3: Self-supervised Learning with Masked Autoencoders (MAE) 15 | - Part 4 (Take-home task): Contrastive Language-Image Pre-training (CLIP) 16 | - Part 5: Conclusion 17 | 18 | 19 | ### Notebooks 20 | 21 | Tutorial: [![Open In 22 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/1_vision/vision_tutorial.ipynb) 23 | 24 | Tutorial with solutions: [![Open In 25 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/1_vision/vision_tutorial_solutions.ipynb) 26 | 27 | --- 28 | -------------------------------------------------------------------------------- /1_vision/vision_tutorial.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "id": "8t47OV2_FVeX" 7 | }, 8 | "source": [ 9 | "# EEML'25 Tutorial: Training Vision Transformer (ViT) encoders from scratch\n", 10 | "\n", 11 | "Authors:\n", 12 | "\n", 13 | "- Liliane Momeni (lmomeni@google.com)\n", 14 | "- Nikhil Parthasarathy (nikparth@google.com)\n", 15 | "\n", 16 | "\n", 17 | "**Abstract:** This tutorial is designed to provide a hands-on understanding of how to train powerful Vision Transformer (ViT) encoders from scratch 🚀. Part 1 will recap why ViT encoders are important, before diving into the three fundamental training paradigms. Part 2 will cover the classic fully supervised approach using labeled data. Part 3 will explore the self-supervised paradigm through Masked Autoencoders, where models learn rich features from the data itself without labels. Finally, Part 4 covers contrastive learning with noisy supervision, using CLIP as our case study.\n", 18 | "\n", 19 | "**Outline**:\n", 20 | "- Part 1: Introduction to Vision Transformer encoders\n", 21 | "- Part 2: Fully supervised ViT for Image Classification\n", 22 | "- Part 3: Self-supervised Learning with Masked Autoencoders (MAE)\n", 23 | "- Part 4 (Take-home task): Contrastive Language-Image Pre-training (CLIP)\n", 24 | "- Part 5: Conclusion\n", 25 | "\n", 26 | "**Slides**: The set of slides which accompany this tutorial can be found [here](https://docs.google.com/presentation/d/1TeG1pd0Bx4ijSEAFJNdUBGmm9tAvvbwZg8URty7cjWQ/edit?usp=sharing).\n", 27 | "\n", 28 | "❗ Note: Parts 2, 3 and 4 of this tutorial do not require a GPU and can be run on a CPU. If you are able to obtain a Colab runtime with a GPU (you can do this by clicking Runtime -> Change runtime type, and set the hardware accelerator to GPU), it will run slightly faster.\n" 29 | ] 30 | }, 31 | { 32 | "cell_type": "markdown", 33 | "metadata": { 34 | "id": "ncuPRuFWGaxC" 35 | }, 36 | "source": [ 37 | "## Part 1: Introduction to Vision Transformer encoders\n", 38 | "\n", 39 | "**1.1 Vision Encoders**\n", 40 | "\n", 41 | "A vision encoder is a neural network that converts an image into a compact numerical summary called an embedding or feature vector.\n", 42 | "\n", 43 | "🖼️➡️🔢\n", 44 | "\n", 45 | "- **What it does**: It acts as a translator, distilling complex pixel data into a format that a machine can understand and analyze.\n", 46 | "\n", 47 | "- **What it captures**: This summary contains the image's essential visual and semantic content, such as the objects within it, their geometry, and how they relate to each other in the scene.\n", 48 | "\n", 49 | "- **Why it's important**: By learning to create high-quality representations, vision encoders enable models to perform a vast array of visual tasks (such as image classification, object detection, semantic segmentation, and even image generation) with remarkable accuracy.\n", 50 | "\n", 51 | "**1.2 The Vision Transformer (ViT): A Paradigm Shift**\n", 52 | "\n", 53 | "As the Transformers architecture scaled well in Natural Language Processing, the same architecture was applied to images by creating small 2D patches of the image and treating them as tokens. The result was a Vision Transformer (ViT).\n", 54 | "\n", 55 | "The primary strength of the ViT architecture is its ability to capture **global context** and **long-range dependencies** from the very first layer, a key difference from the hierarchical, localized feature extraction of CNNs. Summarising key differences below:\n", 56 | "\n", 57 | "- Convolutional Neural Networks (CNNs): Process images hierarchically. They use sliding filters (kernels) to capture local features like edges and textures first. They must pass through many layers to gradually build up a global understanding of the entire scene.\n", 58 | "\n", 59 | "- Vision Transformers (ViTs): Take a global-first approach. By treating an image as a sequence of patches, they use a mechanism known as self-attention (which we will explain in Part 2) to immediately understand how a patch in the top-left corner relates to a patch in the bottom-right, leading to a more holistic initial understanding of the scene.\n", 60 | "\n", 61 | "**1.3 The Standard Workflow: Using Pretrained Models**\n", 62 | "\n", 63 | "In practice, training a large model like a Vision Transformer (ViT) from scratch is computationally expensive and therefore rare. The standard workflow is to use a pretrained model, which has already learned a robust, general understanding of the visual world from a massive dataset like ImageNet-21k (containing over 14 million images).\n", 64 | "\n", 65 | "Once you have this strong foundation, you typically follow one of two paths:\n", 66 | "\n", 67 | "- **Fine-Tuning**: You adapt the model by continuing its training on a smaller, task-specific dataset (e.g., classifying medical images or identifying product defects).\n", 68 | "\n", 69 | "- **Zero-Shot Inference**: For certain models, you can use them directly to make predictions on tasks or classes they have never explicitly been trained on.\n", 70 | "\n", 71 | "While using pretrained models is the common approach, this tutorial will pull back the curtain and show you how these powerful encoders are built and trained from the ground up, giving you a fundamental understanding of how they truly work.\n", 72 | "\n", 73 | "**1.4 Different training paradigms**\n", 74 | "\n", 75 | "\n", 76 | "There are three primary paradigms for training these encoders, each with its own philosophy for learning visual representations. We will explore all of them in this tutorial series:\n", 77 | "\n", 78 | "- **Fully Supervised Learning**: This is the classic approach where the model learns from a dataset containing images and their corresponding explicit labels (e.g., a picture of a golden retriever with the label \"dog\"). The model's goal is to correctly predict the label for a given image.\n", 79 | "\n", 80 | "- **Self-Supervised Learning (SSL)**: This powerful paradigm enables a model to learn rich features directly from unlabeled data by solving a \"pretext task\". We will implement a Masked Autoencoder (MAE), where the model must reconstruct randomly hidden patches of an image, forcing it to learn the context and structure of the visual world.\n", 81 | "\n", 82 | "- **Multimodal Contrastive Learning**: This approach trains a model to understand the relationship between data by learning which samples are similar and which are different. We will use CLIP (Contrastive Language-Image Pre-Training), which learns to match images with their corresponding text descriptions scraped from the web. The text provides a \"noisy\" but highly effective form of supervision.\n", 83 | "\n", 84 | "\n" 85 | ] 86 | }, 87 | { 88 | "cell_type": "markdown", 89 | "metadata": { 90 | "id": "8nJjcw0NoQDK" 91 | }, 92 | "source": [ 93 | "## Part 2: Fully supervised [ViT](https://arxiv.org/abs/2010.11929) for Image Classification\n", 94 | "\n", 95 | "This section will walk you through the fundamental components and training process for a Vision Transformer (ViT) in a fully supervised setting.\n", 96 | "We will cover everything from loss function implementation and data preparation to model training and the critical analysis of its performance.\n", 97 | "Each subsection corresponds to a hands-on exercise.\n", 98 | "\n", 99 | "First off run the next cell to install the necessary libraries. As it might take a bit of time to finish installing, start reading Part 2.1." 100 | ] 101 | }, 102 | { 103 | "cell_type": "code", 104 | "execution_count": null, 105 | "metadata": { 106 | "id": "wdG_pA1w4d9Q" 107 | }, 108 | "outputs": [], 109 | "source": [ 110 | "# =========================================\n", 111 | "# Setup and Imports\n", 112 | "# =========================================\n", 113 | "# First, let's install the necessary libraries.\n", 114 | "!pip install -q torch torchvision timm matplotlib\n", 115 | "\n", 116 | "import torch\n", 117 | "import torch.nn as nn\n", 118 | "import torch.optim as optim\n", 119 | "import torch.nn.functional as F\n", 120 | "from torchvision import datasets, transforms, models\n", 121 | "from torch.utils.data import DataLoader\n", 122 | "import torchvision\n", 123 | "import timm\n", 124 | "import math\n", 125 | "\n", 126 | "import numpy as np\n", 127 | "import matplotlib.pyplot as plt\n", 128 | "from PIL import Image\n", 129 | "import requests\n", 130 | "from tqdm.notebook import tqdm\n", 131 | "\n", 132 | "print(f\"PyTorch version: {torch.__version__}\")\n", 133 | "print(f\"Timm version: {timm.__version__}\")\n", 134 | "\n", 135 | "# Set device\n", 136 | "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", 137 | "print(f\"Using device: {device}\")" 138 | ] 139 | }, 140 | { 141 | "cell_type": "markdown", 142 | "metadata": { 143 | "id": "xhpn3hRhpUEC" 144 | }, 145 | "source": [ 146 | "### Part 2.1: The Cross-Entropy Loss Function\n", 147 | "\n", 148 | "At the heart of any classification task is the loss function, which quantifies how \"wrong\" our model's predictions are compared to the actual labels. For multi-class classification, the standard choice is the **Cross-Entropy Loss**.\n", 149 | "\n", 150 | "Before calculating the loss, a model outputs raw, unbounded scores called logits. To turn these logits into a probability distribution that sums to 1, we use the **Softmax** function. Below, we show an example where we are trying to classify 3 images (of size 32x32 pixels) and the total number of classes $C$ is equal to 3 (cat, dog, horse).\n", 151 | "\n", 152 | "
\n", 153 | "
\n", 154 | "\n", 155 | "\n", 156 | "The Softmax equation for the score of a class $i$ is:\n", 157 | "\n", 158 | "$$ \\text{Softmax}(z_i) = \\frac{e^{z_i}}{\\sum_{j=1}^{C} e^{z_j}} $$\n", 159 | "\n", 160 | "Where:\n", 161 | "\n", 162 | "* $z_i$ is the logit for class $i$.\n", 163 | "* $C$ is the total number of classes.\n", 164 | "\n", 165 | "\n", 166 | "Cross-entropy then measures the difference between this predicted probability distribution and the \"true\" distribution (where the correct class has a probability of 1 and all others have a probability of 0).\n", 167 | "\n", 168 | "Mathematically, for a single example, the cross-entropy loss is:\n", 169 | "\n", 170 | "$$L_{CE} = - \\sum_{i=1}^{C} y_i \\log(\\hat{y}_i)$$\n", 171 | "\n", 172 | "Where:\n", 173 | "\n", 174 | "* $C$ is the number of classes.\n", 175 | "* $y_i$ is the true probability for class $i$ (it's 1 for the correct class, 0 otherwise).\n", 176 | "* $\\hat{y}_i$ is the model's predicted probability for class $i$.\n", 177 | "\n", 178 | "When you're training a classification model and using cross-entropy as your loss function, you are effectively trying to minimize the negative log-likelihood of your model's predictions with respect to the true labels.\n", 179 | "Note, below we plot the range of the **Negative Log-Likelihood** operation: $L(\\textbf{y}) = - \\log(\\textbf{y})$\n", 180 | "\n", 181 | "
\n", 182 | "
\n", 183 | "\n", 184 | "We observe the loss function heavily penalizes the model if it assigns a low probability to the correct class (reaching infinity when the predicted probability is 0), while no loss is applied when the predicted probability is 1. \n", 185 | "\n", 186 | "
\n", 187 | "\n", 188 | "
\n", 189 | "\n", 190 | "**Take-away**: The Cross-Entropy loss combines two steps: 1. A softmax function to convert the model's raw output scores (logits) into probabilities. 2. It then calculates the Negative Log-Likelihood, which heavilily penalises the model if it assigns a low probability to the correct class.\n", 191 | "\n", 192 | "In the next exercise, you will implement this from scratch to solidify your understanding before we use PyTorch's built-in version.\n", 193 | "\n" 194 | ] 195 | }, 196 | { 197 | "cell_type": "code", 198 | "execution_count": null, 199 | "metadata": { 200 | "id": "MnKG6wf8DNmp" 201 | }, 202 | "outputs": [], 203 | "source": [ 204 | "#@title Hint 1 { display-mode: \"form\" }\n", 205 | "\n", 206 | "# The task is to implement the Cross-Entropy loss.\n", 207 | "# Remember the first step is Softmax.\n", 208 | "# Softmax takes logits and converts them to probabilities. You can make use of torch.exp() for this step.\n", 209 | "\n", 210 | "\n", 211 | "# Step 1: Solution\n", 212 | "# exps = torch.exp(logits)\n", 213 | "# sum_exps = torch.sum(exps, dim=1, keepdim=True)\n", 214 | "# softmax_probs = exps / sum_exps\n", 215 | "\n", 216 | "\n", 217 | "# Example\n", 218 | "B, C = 3, 3 # Batch size x Classes\n", 219 | "logits = torch.randn(size=(B, C)) # Batch size x Classes\n", 220 | "exps = torch.exp(logits)\n", 221 | "sum_exps = torch.sum(exps, dim=1, keepdim=True)\n", 222 | "softmax_probs = exps / sum_exps\n", 223 | "\n", 224 | "print(softmax_probs)\n", 225 | "assert torch.allclose(softmax_probs.sum(dim=1), torch.ones(3)) # Make sure they sum to one for each sample" 226 | ] 227 | }, 228 | { 229 | "cell_type": "code", 230 | "execution_count": null, 231 | "metadata": { 232 | "id": "CHYS9iZsDcRP" 233 | }, 234 | "outputs": [], 235 | "source": [ 236 | "#@title Hint 2 { display-mode: \"form\" }\n", 237 | "\n", 238 | "# The next step is Negative Log-Likelihood.\n", 239 | "# Negative Log-Likelihood takes the predicted probability of the *correct* class and calculates -log(probability).\n", 240 | "# For this step you can make use of torch.log().\n", 241 | "\n", 242 | "\n", 243 | "# Step 2: Solution\n", 244 | "# log_probs = - torch.log(softmax_probs)\n", 245 | "\n", 246 | "# Example\n", 247 | "labels = torch.tensor([0, 2, 1])\n", 248 | "labels_ohe = F.one_hot(labels, num_classes=C) # One-hot encoding of the labels\n", 249 | "log_probs = - torch.log(softmax_probs) * labels_ohe\n", 250 | "print(log_probs)" 251 | ] 252 | }, 253 | { 254 | "cell_type": "code", 255 | "execution_count": null, 256 | "metadata": { 257 | "id": "4c56YaxdD2ze" 258 | }, 259 | "outputs": [], 260 | "source": [ 261 | "#@title Hint 3 { display-mode: \"form\" }\n", 262 | "\n", 263 | "# We then need to select the log-probabilities corresponding to the correct class for each item in the batch.\n", 264 | "# torch.gather is great for this, or you can use integer array indexing.\n", 265 | "\n", 266 | "# Step 3: Solution\n", 267 | "# labels = labels.unsqueeze(1)\n", 268 | "# correct_log_probs = torch.gather(log_probs, 1, labels)\n", 269 | "\n", 270 | "# Example\n", 271 | "correct_log_probs = log_probs.max(dim=1).values\n", 272 | "print(log_probs.max(dim=1).values)" 273 | ] 274 | }, 275 | { 276 | "cell_type": "markdown", 277 | "metadata": { 278 | "id": "D4fsbqtJC_vY" 279 | }, 280 | "source": [] 281 | }, 282 | { 283 | "cell_type": "code", 284 | "execution_count": null, 285 | "metadata": { 286 | "id": "VhbG5H8AEAZE" 287 | }, 288 | "outputs": [], 289 | "source": [ 290 | "#@title Hint 4 { display-mode: \"form\" }\n", 291 | "\n", 292 | "# For multiple examples, you'll average the loss across the batch.\n", 293 | "# For this, we can make use of torch.mean().\n", 294 | "\n", 295 | "# Step 4: Solution\n", 296 | "# loss = torch.mean(correct_log_probs)\n", 297 | "\n", 298 | "# Example\n", 299 | "loss = torch.mean(correct_log_probs)\n", 300 | "print(loss)" 301 | ] 302 | }, 303 | { 304 | "cell_type": "code", 305 | "execution_count": null, 306 | "metadata": { 307 | "id": "ERkq1s-usATB" 308 | }, 309 | "outputs": [], 310 | "source": [ 311 | "# ==============================================================================\n", 312 | "#\n", 313 | "# Implement the Cross-Entropy Loss Function from Scratch\n", 314 | "#\n", 315 | "# ==============================================================================\n", 316 | "\n", 317 | "def cross_entropy_loss_scratch(logits, labels):\n", 318 | " \"\"\"\n", 319 | " Cross-entropy loss for classification.\n", 320 | " Formula: -log(exp(x_i) / sum(exp(x_j))) for the correct class i\n", 321 | "\n", 322 | " Args:\n", 323 | " logits (torch.Tensor): The raw, unnormalized scores from the model.\n", 324 | " Shape: (batch_size, num_classes).\n", 325 | " labels (torch.Tensor): The true class indices. Shape: (batch_size,).\n", 326 | "\n", 327 | " Returns:\n", 328 | " torch.Tensor: A scalar tensor representing the mean loss.\n", 329 | " \"\"\"\n", 330 | "\n", 331 | " ## --- TODO 1: YOUR CODE HERE ---\n", 332 | "\n", 333 | "\n", 334 | " # --- END YOUR CODE --- #\n", 335 | " return loss\n", 336 | "\n", 337 | "\n", 338 | "# --- Verification ---\n", 339 | "print(\"--- Verifying Custom Cross-Entropy Loss ---\")\n", 340 | "# Create some dummy data\n", 341 | "dummy_logits = torch.tensor([[2.0, 1.0, 0.1], [0.5, 2.5, 0.1]], device=device)\n", 342 | "dummy_labels = torch.tensor([0, 1], device=device)\n", 343 | "\n", 344 | "# Calculate loss with our function\n", 345 | "loss_scratch = cross_entropy_loss_scratch(dummy_logits, dummy_labels)\n", 346 | "\n", 347 | "# Calculate loss with PyTorch's built-in function\n", 348 | "loss_pytorch = nn.CrossEntropyLoss()(dummy_logits, dummy_labels)\n", 349 | "\n", 350 | "print(f\"Loss from our implementation: {loss_scratch.item():.4f}\")\n", 351 | "print(f\"Loss from PyTorch's nn.CrossEntropyLoss: {loss_pytorch.item():.4f}\")\n", 352 | "# The values should be very close!\n", 353 | "assert torch.allclose(loss_scratch, loss_pytorch), \"Our implementation is not correct!\"\n", 354 | "print(\"-\" * 40, \"\\n\")\n" 355 | ] 356 | }, 357 | { 358 | "cell_type": "markdown", 359 | "metadata": { 360 | "id": "esaTzeAEE-Q5" 361 | }, 362 | "source": [ 363 | "### Part 2.2: Tokenizing an Image for a ViT\n", 364 | "\n", 365 | "Transformers, which were originally designed for text, expect a sequence of \\\"tokens\\\" as input. How do we create a sequence from an image? The Vision Transformer's clever solution is to break the image down into a grid of smaller patches.\n", 366 | "\n", 367 | "Process:\n", 368 | "\n", 369 | "1. **Divide**: An input image is split into a grid of fixed-size, non-overlapping patches. For example, if an image of size 224x224 pixels is split into patches of size 16x16 pixels, this means there are a total of 14*14 patches.\n", 370 | "\n", 371 | "
\n", 372 | "\n", 373 | "
\n", 374 | "\n", 375 | "2. **Flatten**: Each patch is flattened into a single vector. For a 16x16 patch with 3 color channels (RGB), this results in a vector of size 16 * 16 * 3 = 768.\n", 376 | "\n", 377 | "3. **Project**: Each of these vectors is then linearly projected into a consistent embedding dimension (D) that the Transformer expects. This is essentially a standard linear layer applied to each patch vector.\n", 378 | "\n", 379 | "This sequence of patch embeddings is the input for our Transformer.\n", 380 | "\n", 381 | "In practice, Steps 1, 2 and 3 can be done efficiently by simply using a 2D convolutional layer where the kernel size and stride are equal to the patch size.\n", 382 | "\n" 383 | ] 384 | }, 385 | { 386 | "cell_type": "markdown", 387 | "metadata": { 388 | "id": "T_cKMhKVMlfC" 389 | }, 390 | "source": [ 391 | "### Part 2.3: Positional Embeddings\n", 392 | "\n", 393 | "The patching process creates a sequence of tokens, but it discards all spatial information. The Transformer has no inherent sense of which patch came from the top-left corner versus the bottom-right.\n", 394 | "\n", 395 | "To solve this, we add **Positional Embeddings** to the patch embeddings. These are learnable vectors that encode the position of each patch in the original image grid. This gives the model crucial information about the original location of each patch, allowing it to understand the image's structure.\n", 396 | "\n", 397 | "Think of it like this: if you tear a photo into pieces and shuffle them, you lose the picture. Adding positional embeddings is like writing the original coordinate (e.g. 'top-left') on the back of each piece, allowing the model to understand the spatial relationships.\n", 398 | "\n", 399 | "**Take-away**, the steps to convert an image into token embeddings to feed to the Vision Transformer include:\n", 400 | "- Tokenize the image with a Conv2D or by separately (i) dividing the image into patches (ii) flattening the patches and (iii) linear projecting the patch embeddings\n", 401 | "- Sum patch embeddings and learned positional embeddings\n", 402 | "\n", 403 | "
\n", 404 | "\n", 405 | "
" 406 | ] 407 | }, 408 | { 409 | "cell_type": "code", 410 | "execution_count": null, 411 | "metadata": { 412 | "id": "L2DmJ7NOQNAl" 413 | }, 414 | "outputs": [], 415 | "source": [ 416 | "#@title Hint 1 { display-mode: \"form\" }\n", 417 | "\n", 418 | "# The first step is to project the image into patch embeddings.\n", 419 | "# A convolutional layer is a great way to do this efficiently.\n", 420 | "# A kernel size and stride equal to `patch_size` will make the conv\n", 421 | "# layer act on non-overlapping patches.\n", 422 | "\n", 423 | "# Step 1: Solution\n", 424 | "# x = self.proj(x)\n", 425 | "\n", 426 | "# Example\n", 427 | "patch_size = 16\n", 428 | "in_chans = 3\n", 429 | "embed_dim = 768\n", 430 | "proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n", 431 | "x = torch.randn(1, 3, 224, 224) # Dummy image tensor\n", 432 | "print(f\"Shape before projection: {x.shape}\") # Shape: (Batch, Channels, Height, Width)\n", 433 | "x = proj(x)\n", 434 | "print(f\"Shape after projection: {x.shape}\") # Shape: (Batch, embed_dim, H/patch_size, W/patch_size)" 435 | ] 436 | }, 437 | { 438 | "cell_type": "code", 439 | "execution_count": null, 440 | "metadata": { 441 | "id": "EGBkckUbTK9l" 442 | }, 443 | "outputs": [], 444 | "source": [ 445 | "#@title Hint 2 { display-mode: \"form\" }\n", 446 | "\n", 447 | "# What is the shape of the output of the previous step? We need to flatten the spatial dimensions.\n", 448 | "\n", 449 | "# Step 2: Solution\n", 450 | "# x = x.flatten(2)\n", 451 | "\n", 452 | "# Example (continuing with the projected x from the previous hint)\n", 453 | "print(f\"Shape before flattening: {x.shape}\") # Shape: (Batch, embed_dim, H/patch_size, W/patch_size)\n", 454 | "x = x.flatten(2)\n", 455 | "print(f\"Shape after flattening: {x.shape}\") # Shape: (Batch, embed_dim, N_patches)" 456 | ] 457 | }, 458 | { 459 | "cell_type": "code", 460 | "execution_count": null, 461 | "metadata": { 462 | "id": "OZCkRAxTZc_2" 463 | }, 464 | "outputs": [], 465 | "source": [ 466 | "#@title Hint 3 { display-mode: \"form\" }\n", 467 | "\n", 468 | "# Now we need to add the positional embeddings. Be careful with the shapes!\n", 469 | "\n", 470 | "# Step 3: Solution\n", 471 | "# x = x.transpose(1, 2)\n", 472 | "# x = x + self.pos_embed\n", 473 | "\n", 474 | "# Example (continuing with flattened x from the previous hint)\n", 475 | "print(f\"Shape before transpose: {x.shape}\") # Shape: (Batch, embed_dim, N_patches)\n", 476 | "x = x.transpose(1, 2)\n", 477 | "print(f\"Shape of x after transpose: {x.shape}\") # Shape: (Batch, N_patches, embed_dim)\n", 478 | "pos_embed = torch.randn(1, 196, 768)\n", 479 | "print(f\"Shape of pos_embed: {pos_embed.shape}\") # Shape: (Batch, N_patches, embed_dim)\n", 480 | "x = x + pos_embed\n", 481 | "print(f\"Shape of x after adding pos_embed: {x.shape}\") # Shape: (1, N_patches, embed_dim)" 482 | ] 483 | }, 484 | { 485 | "cell_type": "code", 486 | "execution_count": null, 487 | "metadata": { 488 | "id": "EmJYNAa6PFL9" 489 | }, 490 | "outputs": [], 491 | "source": [ 492 | "# ==============================================================================\n", 493 | "#\n", 494 | "# Image Tokenization and Positional Embeddings\n", 495 | "#\n", 496 | "# ==============================================================================\n", 497 | "\n", 498 | "class PatchAndEmbed(nn.Module):\n", 499 | " \"\"\"\n", 500 | " Takes an image, splits it into patches, and embeds them.\n", 501 | " Also adds positional embeddings.\n", 502 | " \"\"\"\n", 503 | " def __init__(self, img_size, patch_size, in_chans, embed_dim):\n", 504 | " super().__init__()\n", 505 | " self.img_size = img_size\n", 506 | " self.patch_size = patch_size\n", 507 | " self.n_patches = (img_size // patch_size) ** 2\n", 508 | "\n", 509 | " # 1. Create the projection layer\n", 510 | " self.proj = nn.Conv2d(\n", 511 | " in_chans,\n", 512 | " embed_dim,\n", 513 | " kernel_size=patch_size,\n", 514 | " stride=patch_size\n", 515 | " )\n", 516 | "\n", 517 | " # 2. Create the learnable positional embeddings.\n", 518 | " self.pos_embed = nn.Parameter(torch.zeros(1, self.n_patches, embed_dim))\n", 519 | "\n", 520 | "\n", 521 | " def forward(self, x):\n", 522 | " \"\"\"\n", 523 | " Args:\n", 524 | " x (torch.Tensor): Input images. Shape: (batch_size, in_chans, img_size, img_size)\n", 525 | " Returns:\n", 526 | " torch.Tensor: Embedded patches. Shape: (batch_size, n_patches, embed_dim)\n", 527 | " \"\"\"\n", 528 | "\n", 529 | " ## --- TODO 1: YOUR CODE HERE ---\n", 530 | "\n", 531 | "\n", 532 | " # --- END YOUR CODE --- #\n", 533 | " return x\n", 534 | "\n", 535 | "\n", 536 | "# --- Verification ---\n", 537 | "print(\"--- Verifying Patch and Embed Layer ---\")\n", 538 | "patcher = PatchAndEmbed(img_size=224, patch_size=16, in_chans=3, embed_dim=768).to(device)\n", 539 | "dummy_image_batch = torch.randn(4, 3, 224, 224).to(device) # Batch of 4 images\n", 540 | "embedded_patches = patcher(dummy_image_batch)\n", 541 | "print(f\"Shape of output tensor: {embedded_patches.shape}\")\n", 542 | "print(\"Expected shape: (4, 196, 768) -> (batch_size, n_patches, embed_dim)\")\n", 543 | "print(\"-\" * 40, \"\\n\")" 544 | ] 545 | }, 546 | { 547 | "cell_type": "markdown", 548 | "metadata": { 549 | "id": "3-1deoktnhkZ" 550 | }, 551 | "source": [ 552 | "### Part 2.4: The ViT Encoder Block\n", 553 | "\n", 554 | "The core of the ViT is a stack of **Encoder Blocks**. Each block has two main sub-layers:\n", 555 | "\n", 556 | "1. **Self-Attention (SA)**\n", 557 | "\n", 558 | "- This is where the magic happens. For each patch in the input sequence, self-attention computes an updated representation by incorporating information from the entire set of patches, putting more weight on relevant patches. In practice, an attention mechanism has three parts:\n", 559 | "\n", 560 | " - **Query (Q)**: The query is a feature vector that represents the current patch's \"search query.\" It asks, \"What kind of information should I look for in other patches to better understand myself?\"\n", 561 | "\n", 562 | " - **Key-Value Pairs (K-V)**: The key is a feature vector that represents each patch's \"label\" or \"identifier.\" It responds to queries by saying, \"This is the kind of information I contain.\" The value is a feature vector that represents the actual content or features of each patch. This is the information that gets passed along if a patch is \"attended to.\"\n", 563 | "\n", 564 | " - **Score function (V)**: The score function takes the query and a key as input and outputs the score / attention weight of the query-key pair. This is usually implemented by a simple similarity metric like a dot product.\n", 565 | "\n", 566 | " Note that the Query and Key-Value Pairs are obtained by multiplying the input with different weight matrics as below:\n", 567 | "\n", 568 | "
\n", 569 | " \n", 570 | "
\n", 571 | "\n", 572 | "- In practice, the attention mechanism works as follows:\n", 573 | "\n", 574 | " - **Calculate Attention Scores**: For a single patch, its Query vector is compared with the Key vectors of all other patches in the sequence using a dot product. A high score $s$ indicates a strong relevance between the Query and a Key.\n", 575 | "\n", 576 | " $$ s = {QK^{T}} $$\n", 577 | "\n", 578 | " - **Get Attention Weights**: The scores are scaled by the square root of the dimension of the key vectors (to ensure stable gradients) and then passed through a softmax function. This converts the scores into a probability distribution, where the weights sum to 1. These weights $ \\alpha$ determine how much \"attention\" the current patch should pay to every other patch.\n", 579 | "\n", 580 | " $$\\alpha = \\text{softmax}(\\frac{s}{\\sqrt{d_k}})$$\n", 581 | "\n", 582 | " - **Compute Output**: The final output for the patch is a weighted sum of all the Value vectors in the sequence, using the attention weights calculated in the previous step.\n", 583 | "\n", 584 | " $$ o = \\alpha \\cdot V $$\n", 585 | "\n", 586 | " The overall operation is concisely captured by the Scaled Dot-Product Attention formula:\n", 587 | "\n", 588 | " $$\n", 589 | "\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V\n", 590 | "$$\n", 591 | "\n", 592 | "
\n", 593 | " \n", 594 | "
\n", 595 | "\n", 596 | "2. **Multilayer Perceptron (MLP)**\n", 597 | "\n", 598 | "- A standard feed-forward neural network then processes each patch's representation independently after the attention step. This provides additional computational capacity.\n", 599 | "\n", 600 | "**Take-away**:\n", 601 | "- Self-Attention: Its primary role is to weigh the significance of all other tokens in a sequence for a given token, creating a contextually-aware representation by capturing long-range dependencies and relationships within the input data.\n", 602 | "\n", 603 | "- MLP (Multi-Layer Perceptron): This component processes each token's representation from the self-attention layer in isolation, applying a non-linear transformation to enrich the features of each token individually and enhance the model's expressive power." 604 | ] 605 | }, 606 | { 607 | "cell_type": "markdown", 608 | "metadata": { 609 | "id": "h78vl7-dWxZe" 610 | }, 611 | "source": [ 612 | "**Remark**: We usually apply self-attention multiple (*h*) times using different Key (K), Query (Q), and Value (V) matrices. This approach, known as Multi-Head Self-Attention (MHSA), allows the model to learn various relationships between tokens or patches, leading to a richer representation. Each \"head\" can focus on different aspects of the input, and their combined (*concatenated*) outputs provide a more comprehensive understanding.\n", 613 | "\n", 614 | "
\n", 615 | " \n", 616 | "
" 617 | ] 618 | }, 619 | { 620 | "cell_type": "code", 621 | "execution_count": null, 622 | "metadata": { 623 | "id": "zXPi1sr3bEFY" 624 | }, 625 | "outputs": [], 626 | "source": [ 627 | "#@title Hint 1 { display-mode: \"form\" }\n", 628 | "\n", 629 | "# An MLP consists of 3 components:\n", 630 | "# 1. An hidden layer\n", 631 | "# 2. An activation function that follows it. This function introduces non-linearity, allowing the network to learn more complex patterns than a simple linear model.\n", 632 | "# 3. An output layer\n", 633 | "\n", 634 | "\n", 635 | "# Step 1: Solution\n", 636 | "class MLP(nn.Module):\n", 637 | " def __init__(self, in_features, hidden_features, out_features):\n", 638 | " super().__init__()\n", 639 | " self.fc1 = nn.Linear(in_features, hidden_features)\n", 640 | " self.act = nn.GELU()\n", 641 | " self.fc2 = nn.Linear(hidden_features, out_features)\n", 642 | "\n", 643 | " def forward(self, x):\n", 644 | " x = self.fc1(x)\n", 645 | " x = self.act(x)\n", 646 | " x = self.fc2(x)\n", 647 | " return x\n", 648 | "\n", 649 | "\n", 650 | "# Example\n", 651 | "in_features = 768\n", 652 | "hidden_features = 3072 # A common ratio is 4x the input features\n", 653 | "out_features = 768\n", 654 | "\n", 655 | "mlp = MLP(in_features, hidden_features, out_features)\n", 656 | "input_tensor = torch.randn(1, 196, in_features) # (Batch, N_patches, in_features)\n", 657 | "output = mlp(input_tensor)\n", 658 | "print(f\"MLP input shape: {input_tensor.shape}\")\n", 659 | "print(f\"MLP output shape: {output.shape}\")" 660 | ] 661 | }, 662 | { 663 | "cell_type": "code", 664 | "execution_count": null, 665 | "metadata": { 666 | "id": "1JEZvBhFTJ_l" 667 | }, 668 | "outputs": [], 669 | "source": [ 670 | "#@title Hint 2 { display-mode: \"form\" }\n", 671 | "\n", 672 | "# The first step for the scaled dot product attention is to do the dot product of the keys and values.\n", 673 | "\n", 674 | "\n", 675 | "# Step 2: Solution\n", 676 | "# Matmul Q and K transpose: (B, N, C) @ (B, C, N) -> (B, N, N)\n", 677 | "# attn_logits = torch.matmul(q, k.transpose(-2, -1))\n", 678 | "\n", 679 | "\n", 680 | "# Example\n", 681 | "B, N, C = 4, 196, 768 # (batch_size, num_patches, embed_dim)\n", 682 | "q = torch.randn(B, N, C)\n", 683 | "k = torch.randn(B, N, C)\n", 684 | "print(f\"Shape of Query (Q): {q.shape}\")\n", 685 | "print(f\"Shape of Key (K): {k.shape}\")\n", 686 | "\n", 687 | "k_t = k.transpose(-2, -1)\n", 688 | "print(f\"Shape of K transposed: {k_t.shape}\")\n", 689 | "attn_logits = torch.matmul(q, k_t)\n", 690 | "print(f\"Shape of attention logits: {attn_logits.shape}\")" 691 | ] 692 | }, 693 | { 694 | "cell_type": "code", 695 | "execution_count": null, 696 | "metadata": { 697 | "id": "v3dlZSgATnWh" 698 | }, 699 | "outputs": [], 700 | "source": [ 701 | "#@title Hint 3 { display-mode: \"form\" }\n", 702 | "\n", 703 | "# Then we must apply the softmax to get attention weights.\n", 704 | "\n", 705 | "# Step 3: Solution\n", 706 | "# attn_weights = F.softmax(attn_logits, dim=-1)\n", 707 | "\n", 708 | "# Example (continuing from attn_logits from the previous)\n", 709 | "\n", 710 | "attn_weights = F.softmax(attn_logits, dim=-1)\n", 711 | "print(f\"Shape of attention weights: {attn_weights.shape}\")\n", 712 | "\n", 713 | "# The sum of weights for each query should be 1\n", 714 | "print(f\"Sum of weights for the first query in the first batch item: {attn_weights[0, 0, :].sum()}\")" 715 | ] 716 | }, 717 | { 718 | "cell_type": "code", 719 | "execution_count": null, 720 | "metadata": { 721 | "id": "IiPVRny_Ttew" 722 | }, 723 | "outputs": [], 724 | "source": [ 725 | "#@title Hint 4 { display-mode: \"form\" }\n", 726 | "\n", 727 | "# Then we need to multiply the attention weights by the values to get our weighted average.\n", 728 | "\n", 729 | "# Step 4: Solution\n", 730 | "# Matmul attention probabilities with V: (B, N, N) @ (B, N, C) -> (B, N, C)\n", 731 | "# values = torch.matmul(attn_weights, v)\n", 732 | "\n", 733 | "# Example\n", 734 | "B, N, C = 4, 196, 768\n", 735 | "attn_weights = torch.randn(B, N, N)\n", 736 | "v = torch.randn(B, N, C)\n", 737 | "print(f\"Shape of attention weights: {attn_weights.shape}\")\n", 738 | "print(f\"Shape of Values (V): {v.shape}\")\n", 739 | "\n", 740 | "values = torch.matmul(attn_weights, v)\n", 741 | "print(f\"Shape of final values: {values.shape}\")" 742 | ] 743 | }, 744 | { 745 | "cell_type": "code", 746 | "execution_count": null, 747 | "metadata": { 748 | "id": "61-ZVk6-2Tu2" 749 | }, 750 | "outputs": [], 751 | "source": [ 752 | "# ==============================================================================\n", 753 | "#\n", 754 | "# Implement MLP and Self-Attention\n", 755 | "#\n", 756 | "#\n", 757 | "# ==============================================================================\n", 758 | "\n", 759 | "def scaled_dot_product_attention(q, k, v):\n", 760 | " \"\"\"\n", 761 | " Calculates the Scaled Dot-Product Attention.\n", 762 | " The inputs q, k, v are expected to be of shape (Batch, Sequence_Length, Dimension).\n", 763 | " \"\"\"\n", 764 | " d_k = k.size()[-1] # Get the dimension of the key vectors\n", 765 | "\n", 766 | " ## --- TODO 1: YOUR CODE HERE ---\n", 767 | "\n", 768 | " # --- END YOUR CODE --- #\n", 769 | "\n", 770 | " attn_logits = attn_logits / math.sqrt(d_k) # (batch_size, n_patches, n_patches)\n", 771 | "\n", 772 | " ## --- TODO 2: YOUR CODE HERE ---\n", 773 | "\n", 774 | " # --- END YOUR CODE --- #\n", 775 | "\n", 776 | " return values\n", 777 | "\n", 778 | "class SelfAttention(nn.Module):\n", 779 | " def __init__(self, embed_dim):\n", 780 | " \"\"\"\n", 781 | " Args:\n", 782 | " embed_dim (int): The embedding dimension of the input, which is also the\n", 783 | " dimension of Q, K, and V.\n", 784 | " \"\"\"\n", 785 | " super().__init__()\n", 786 | " self.embed_dim = embed_dim\n", 787 | "\n", 788 | " # 1. A single linear layer to project input into Q, K, V for efficiency\n", 789 | " self.qkv = nn.Linear(embed_dim, embed_dim * 3)\n", 790 | " self.proj = nn.Linear(embed_dim, embed_dim)\n", 791 | "\n", 792 | " def forward(self, x):\n", 793 | " # Input x shape: (batch_size, n_patches, embed_dim)\n", 794 | " B, N, C = x.shape\n", 795 | "\n", 796 | " # 2. Project to Q, K, V\n", 797 | " qkv = self.qkv(x) # (batch_size, n_patches, 3 * embed_dim)\n", 798 | "\n", 799 | " # 3. Split the last dimension into 3 chunks for Q, K, V\n", 800 | " q, k, v = qkv.chunk(3, dim=-1) # (batch_size, n_patches, embed_dim)\n", 801 | "\n", 802 | " # 4. Call the scaled dot-product attention function\n", 803 | " x = scaled_dot_product_attention(q, k, v) # (batch_size, n_patches, embed_dim)\n", 804 | "\n", 805 | " # 5. Final linear projection\n", 806 | " x = self.proj(x) # (batch_size, n_patches, embed_dim)\n", 807 | " return x\n", 808 | "\n", 809 | "\n", 810 | "class MLP(nn.Module):\n", 811 | " def __init__(self, in_features, hidden_features, out_features):\n", 812 | " \"\"\"\n", 813 | " Args:\n", 814 | " in_features (int): The number of features in the input tensor.\n", 815 | " hidden_features (int): The number of nodes in the hidden layer.\n", 816 | " out_features (int): The number of features in the output tensor.\n", 817 | " \"\"\"\n", 818 | " super().__init__()\n", 819 | " ## --- TODO 3: YOUR CODE HERE ---\n", 820 | "\n", 821 | " # --- END YOUR CODE --- #\n", 822 | "\n", 823 | " def forward(self, x):\n", 824 | " \"\"\"\n", 825 | " Args:\n", 826 | " x (torch.Tensor): The input tensor.\n", 827 | "\n", 828 | " Returns:\n", 829 | " torch.Tensor: The output tensor.\n", 830 | " \"\"\"\n", 831 | " ## --- TODO 4: YOUR CODE HERE ---\n", 832 | "\n", 833 | " # --- END YOUR CODE --- #\n", 834 | " return x\n", 835 | "\n", 836 | "\n", 837 | "# --- Verification ---\n", 838 | "# Using the final embeddings from the previous task\n", 839 | "# final_embeddings.shape is (4, 196, 768)\n", 840 | "embed_dim = 768\n", 841 | "\n", 842 | "# Attention Block\n", 843 | "attention = SelfAttention(embed_dim=embed_dim).to(device)\n", 844 | "attended_output = attention(embedded_patches)\n", 845 | "print(f\"Attention output shape: {attended_output.shape}\") # Expected: (4, 196, 768)\n", 846 | "\n", 847 | "# MLP Block\n", 848 | "mlp = MLP(in_features=embed_dim, hidden_features=embed_dim * 4, out_features=embed_dim).to(device)\n", 849 | "mlp_output = mlp(attended_output)\n", 850 | "print(f\"MLP output shape: {mlp_output.shape}\") # Expected: (4, 196, 768)\n" 851 | ] 852 | }, 853 | { 854 | "cell_type": "markdown", 855 | "metadata": { 856 | "id": "6A-K5Regziy4" 857 | }, 858 | "source": [ 859 | "### Part 2.5: The Full Forward Pass\n", 860 | "\n", 861 | "Now we assemble all the pieces into a complete Transformer Encoder block and then stack them to create the full ViT model.\n", 862 | "\n", 863 | "**Vision Transformer forward pass**: The forward pass of the VisionTransformer begins by converting the input image into a sequence of flattened patch embeddings using self.patch_embed. This sequence of tokens is then processed sequentially through a series of EncoderBlock layers, which refine the token representations. After the final block, the entire sequence is normalized. To produce a single vector for classification, the model takes the **mean of all token embeddings** across the sequence dimension. This aggregated vector is then passed to the final linear layer (self.head) to generate the classification logits. Note that here we take the mean of all token embeddings for simplicity -- in practice, a [CLS] token is often used to aggregate the sequence information (please refer to [ViT](https://arxiv.org/abs/2010.11929) paper for more details).\n", 864 | "\n", 865 | "
\n", 866 | "\n", 867 | "
\n", 868 | "\n", 869 | "\n", 870 | "Note that in addition to the self attention and MLP described in Part 2.4, there are two additional components in our Encoder Block: normalisation layers and residual connections. These ensure stable training." 871 | ] 872 | }, 873 | { 874 | "cell_type": "code", 875 | "execution_count": null, 876 | "metadata": { 877 | "id": "RB_2-EwJZDIm" 878 | }, 879 | "outputs": [], 880 | "source": [ 881 | "#@title Hint 1 { display-mode: \"form\" }\n", 882 | "\n", 883 | "# Call PatchAndEmbed to tokenize the image and add positional embeddings.\n", 884 | "\n", 885 | "# Step 1: Solution\n", 886 | "# x = self.patch_embed(x)\n", 887 | "\n", 888 | "# Example\n", 889 | "\n", 890 | "patch_embed = PatchAndEmbed(img_size=224, patch_size=16, in_chans=3, embed_dim=768)\n", 891 | "img = torch.randn(1, 3, 224, 224)\n", 892 | "output = patch_embed(img)\n", 893 | "\n", 894 | "print(f\"Input image shape: {img.shape}\")\n", 895 | "print(f\"Output shape after PatchAndEmbed: {output.shape}\")" 896 | ] 897 | }, 898 | { 899 | "cell_type": "code", 900 | "execution_count": null, 901 | "metadata": { 902 | "id": "xAG5ThV4ZNgP" 903 | }, 904 | "outputs": [], 905 | "source": [ 906 | "#@title Hint 2 { display-mode: \"form\" }\n", 907 | "\n", 908 | "# Pass through the Transformer encoder blocks.\n", 909 | "\n", 910 | "# Step 2: Solution\n", 911 | "# for blk in self.blocks:\n", 912 | "# x = blk(x)" 913 | ] 914 | }, 915 | { 916 | "cell_type": "code", 917 | "execution_count": null, 918 | "metadata": { 919 | "id": "9w_p15o2Zdfo" 920 | }, 921 | "outputs": [], 922 | "source": [ 923 | "#@title Hint 3 { display-mode: \"form\" }\n", 924 | "\n", 925 | "# Take the average over the sequence dimension and pass to the classification head.\n", 926 | "\n", 927 | "\n", 928 | "# Step 3: Solution\n", 929 | "# x = x.mean(dim=1)\n", 930 | "# logits = self.head(x)\n", 931 | "\n", 932 | "# Example\n", 933 | "B, N, C = 1, 196, 768 # Batch, Num_Patches, Embed_dim\n", 934 | "x = torch.randn(B, N, C)\n", 935 | "print(f\"Shape after Transformer blocks: {x.shape}\")\n", 936 | "\n", 937 | "x_mean = x.mean(dim=1)\n", 938 | "print(f\"Shape after averaging: {x_mean.shape}\")\n", 939 | "\n", 940 | "num_classes = 10 # Example: CIFAR has 10 classes\n", 941 | "head = nn.Linear(C, num_classes)\n", 942 | "\n", 943 | "logits = head(x_mean)\n", 944 | "print(f\"Shape of final logits: {logits.shape}\")" 945 | ] 946 | }, 947 | { 948 | "cell_type": "code", 949 | "execution_count": null, 950 | "metadata": { 951 | "id": "iOh5_G4VVkHy" 952 | }, 953 | "outputs": [], 954 | "source": [ 955 | "# ==============================================================================\n", 956 | "#\n", 957 | "# Full forward pass\n", 958 | "#\n", 959 | "#\n", 960 | "# ==============================================================================\n", 961 | "\n", 962 | "\n", 963 | "\n", 964 | "class EncoderBlock(nn.Module):\n", 965 | " def __init__(self, embed_dim, mlp_ratio=4.0):\n", 966 | " super().__init__()\n", 967 | " self.norm1 = nn.LayerNorm(embed_dim)\n", 968 | " self.attn = SelfAttention(embed_dim)\n", 969 | " self.norm2 = nn.LayerNorm(embed_dim)\n", 970 | " hidden_features = int(embed_dim * mlp_ratio)\n", 971 | " self.mlp = MLP(in_features=embed_dim, hidden_features=hidden_features, out_features=embed_dim)\n", 972 | "\n", 973 | " def forward(self, x):\n", 974 | " # 1. Residual connection around the Attention block\n", 975 | " x = x + self.attn(self.norm1(x))\n", 976 | " # 2. Residual connection around the MLP block\n", 977 | " x = x + self.mlp(self.norm2(x))\n", 978 | " return x\n", 979 | "\n", 980 | "class VisionTransformer(nn.Module):\n", 981 | " def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=10,\n", 982 | " embed_dim=768, depth=12, mlp_ratio=4.0):\n", 983 | " super().__init__()\n", 984 | "\n", 985 | " # 1. Patch + Positional Embeddings\n", 986 | " self.patch_embed = PatchAndEmbed(img_size, patch_size, in_chans, embed_dim)\n", 987 | "\n", 988 | " # 2. Transformer Encoder\n", 989 | " self.blocks = nn.ModuleList([\n", 990 | " EncoderBlock(embed_dim, mlp_ratio)\n", 991 | " for _ in range(depth)\n", 992 | " ])\n", 993 | "\n", 994 | " # 3. Classification Head\n", 995 | " self.norm = nn.LayerNorm(embed_dim)\n", 996 | " self.head = nn.Linear(embed_dim, num_classes)\n", 997 | "\n", 998 | " def forward(self, x):\n", 999 | "\n", 1000 | " ## --- TODO 1: YOUR CODE HERE ---\n", 1001 | "\n", 1002 | " # --- END YOUR CODE --- #\n", 1003 | "\n", 1004 | " # After the final block, the sequence is normalised\n", 1005 | " x = self.norm(x)\n", 1006 | "\n", 1007 | " ## --- TODO 2: YOUR CODE HERE ---\n", 1008 | "\n", 1009 | " # --- END YOUR CODE --- #\n", 1010 | "\n", 1011 | " return logits\n", 1012 | "\n", 1013 | "# --- Example Usage ---\n", 1014 | "# Instantiate the full ViT model\n", 1015 | "vit_model = VisionTransformer(img_size=32, num_classes=10).to(device) # e.g., for CIFAR-10\n", 1016 | "\n", 1017 | "# A batch of 4 images, 3 channels (RGB), 224x224 pixels\n", 1018 | "images = torch.randn(4, 3, 32, 32, device=device)\n", 1019 | "\n", 1020 | "# Get the model's output (logits)\n", 1021 | "logits = vit_model(images)\n", 1022 | "print(f\"Input image shape: {images.shape}\")\n", 1023 | "print(f\"Output logits shape: {logits.shape}\") # Expected: (4, 10)" 1024 | ] 1025 | }, 1026 | { 1027 | "cell_type": "markdown", 1028 | "metadata": { 1029 | "id": "lA3G7LAyZ14X" 1030 | }, 1031 | "source": [ 1032 | "### Part 2.6: The Training Loop\n", 1033 | "\n", 1034 | "With the model defined, the final step is to train it. This involves feeding it data, calculating the loss, and updating the model's weights using an optimizer. We'll run this for one epoch to see the process in action.\n", 1035 | "\n", 1036 | "We'll use the CIFAR-10 dataset as an example. You can visualise some samples below." 1037 | ] 1038 | }, 1039 | { 1040 | "cell_type": "code", 1041 | "execution_count": null, 1042 | "metadata": { 1043 | "id": "ADmtW08caBKC" 1044 | }, 1045 | "outputs": [], 1046 | "source": [ 1047 | "# ==============================================================================\n", 1048 | "#\n", 1049 | "# CIFAR-10 Dataset\n", 1050 | "#\n", 1051 | "#\n", 1052 | "# ==============================================================================\n", 1053 | "\n", 1054 | "# 1. Setup Dataset and DataLoader\n", 1055 | "transform = transforms.Compose([\n", 1056 | " transforms.ToTensor(),\n", 1057 | " transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n", 1058 | "])\n", 1059 | "\n", 1060 | "# 2. Use CIFAR-10 dataset\n", 1061 | "train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n", 1062 | "test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)\n", 1063 | "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)\n", 1064 | "test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=2)\n", 1065 | "\n", 1066 | "classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')\n", 1067 | "dataiter = iter(train_loader)\n", 1068 | "images, labels = next(dataiter)\n", 1069 | "\n", 1070 | "def imshow(img, title=None):\n", 1071 | " img = img / 2 + 0.5 # unnormalize\n", 1072 | " npimg = img.numpy()\n", 1073 | " plt.imshow(np.transpose(npimg, (1, 2, 0)))\n", 1074 | " if title is not None:\n", 1075 | " plt.title(title)\n", 1076 | " plt.show()\n", 1077 | "\n", 1078 | "print('Example Images:')\n", 1079 | "for i in range(0, 5):\n", 1080 | " imshow(images[i], title=classes[labels[i]])\n" 1081 | ] 1082 | }, 1083 | { 1084 | "cell_type": "markdown", 1085 | "metadata": { 1086 | "id": "-fUs00tOomxZ" 1087 | }, 1088 | "source": [ 1089 | "While you wait for this subsequent training cell to run, you can start reading Part 2.7." 1090 | ] 1091 | }, 1092 | { 1093 | "cell_type": "code", 1094 | "execution_count": null, 1095 | "metadata": { 1096 | "id": "E3MxL7AWaVtK" 1097 | }, 1098 | "outputs": [], 1099 | "source": [ 1100 | "# ==============================================================================\n", 1101 | "#\n", 1102 | "# Run a Training Loop for One Epoch\n", 1103 | "#\n", 1104 | "#\n", 1105 | "# ==============================================================================\n", 1106 | "\n", 1107 | "# Initialize Model, Loss, and Optimizer\n", 1108 | "model = VisionTransformer(img_size=32, patch_size=4, in_chans=3, num_classes=10,\n", 1109 | " embed_dim=64, depth=4, mlp_ratio=2.0).to(device)\n", 1110 | "model = torch.compile(model)\n", 1111 | "criterion = nn.CrossEntropyLoss()\n", 1112 | "optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-4)\n", 1113 | "epochs=2\n", 1114 | "# The Training Loop\n", 1115 | "def train_one_epoch(model, loader, criterion, optimizer, device, epoch):\n", 1116 | " model.train()\n", 1117 | " total_loss = 0\n", 1118 | " correct_predictions = 0\n", 1119 | " total_samples = 0\n", 1120 | " pbar = tqdm(loader, desc=f\"Epoch {epoch+1}/{epochs}\")\n", 1121 | "\n", 1122 | " for images, labels in pbar:\n", 1123 | " images, labels = images.to(device), labels.to(device)\n", 1124 | "\n", 1125 | " # 1. Zero the gradients\n", 1126 | " optimizer.zero_grad()\n", 1127 | "\n", 1128 | " # 2. Forward pass\n", 1129 | " outputs = model(images)\n", 1130 | "\n", 1131 | " # 3. Calculate loss\n", 1132 | " loss = criterion(outputs, labels)\n", 1133 | "\n", 1134 | " # 4. Backward pass and optimization\n", 1135 | " loss.backward()\n", 1136 | " optimizer.step()\n", 1137 | "\n", 1138 | " # 5. Track statistics\n", 1139 | " total_loss += loss.item()\n", 1140 | " pbar.set_postfix({'loss': loss.item()})\n", 1141 | "\n", 1142 | " _, predicted = torch.max(outputs.data, 1)\n", 1143 | " total_samples += labels.size(0)\n", 1144 | " correct_predictions += (predicted == labels).sum().item()\n", 1145 | "\n", 1146 | " avg_loss = total_loss / len(loader)\n", 1147 | " accuracy = 100 * correct_predictions / total_samples\n", 1148 | " return avg_loss, accuracy\n", 1149 | "\n", 1150 | "print(\"\\nStarting training for two epoch...\")\n", 1151 | "for epoch in range(epochs):\n", 1152 | " avg_loss, accuracy = train_one_epoch(model, train_loader, criterion, optimizer, device, epoch)\n", 1153 | " print(f\"\\n--- End of Epoch {epoch}---\")\n", 1154 | " print(f\"Average Training Loss: {avg_loss:.4f}\")\n", 1155 | " print(f\"Training Accuracy: {accuracy:.2f}%\")" 1156 | ] 1157 | }, 1158 | { 1159 | "cell_type": "markdown", 1160 | "metadata": { 1161 | "id": "2edyPY3jsJ-S" 1162 | }, 1163 | "source": [ 1164 | "After a single epoch, the performance is already much better than random chance (10 percent given there are 10 classes)." 1165 | ] 1166 | }, 1167 | { 1168 | "cell_type": "markdown", 1169 | "metadata": { 1170 | "id": "Ck7lHMdVowpf" 1171 | }, 1172 | "source": [ 1173 | "### Part 2.7: Loading Pretrained ViT\n", 1174 | "\n", 1175 | "As mentioned in the introduction, it is common to start from a pretrained model, which has already learned a robust, general understanding of the visual world from a dataset like ImageNet-1k (containing 1 million images). Let's give this a go now and run the pretrained model on a few samples of the CIFAR-10 test set and compare to what we get from evaluating our much smaller model!" 1176 | ] 1177 | }, 1178 | { 1179 | "cell_type": "code", 1180 | "execution_count": null, 1181 | "metadata": { 1182 | "id": "8tS2OMUwqaU1" 1183 | }, 1184 | "outputs": [], 1185 | "source": [ 1186 | "# ==============================================================================\n", 1187 | "#\n", 1188 | "# Inference with a Pretrained ViT\n", 1189 | "#\n", 1190 | "# ==============================================================================\n", 1191 | "from torchvision.transforms import Resize\n", 1192 | "# Load the pretrained Vision Transformer model\n", 1193 | "from transformers import ViTForImageClassification\n", 1194 | "model_pt = ViTForImageClassification.from_pretrained(\"google/vit-base-patch16-224\")\n", 1195 | "model_pt = model_pt.to(device)\n", 1196 | "model_pt.eval() # Set the model to evaluation mode\n", 1197 | "model.eval()\n", 1198 | "# Create a resize transform object\n", 1199 | "resize_transform = Resize((224, 224), antialias=True)\n", 1200 | "\n", 1201 | "print(\"\\n--- Evaluating on 10 samples from the CIFAR-10 test set ---\")\n", 1202 | "\n", 1203 | "# 1. Use torch.no_grad() for inference\n", 1204 | "with torch.no_grad():\n", 1205 | " # 2. Get just one batch of images and labels from the test loader\n", 1206 | " images, labels = next(iter(test_loader))\n", 1207 | "\n", 1208 | " # 3. Resize the entire batch of images to 224x224 to match input size of pretrained model\n", 1209 | " images_rs = resize_transform(images)\n", 1210 | "\n", 1211 | " # 4. Move data to the selected device\n", 1212 | " images_rs = images_rs.to(device)\n", 1213 | " labels = labels.to(device)\n", 1214 | "\n", 1215 | " # 5. Get model predictions (logits)\n", 1216 | " outputs_pt = model_pt(images_rs)\n", 1217 | " logits = outputs_pt.logits\n", 1218 | " outputs_ours = model(images)\n", 1219 | "\n", 1220 | "\n", 1221 | " # Print results for the first 10 images in the batch\n", 1222 | " for i in range(10):\n", 1223 | " true_label = classes[labels[i]]\n", 1224 | " prediction_pt = logits[i].argmax(-1)\n", 1225 | " prediction_ours = outputs_ours.data[i].argmax(-1)\n", 1226 | " imshow(images[i])\n", 1227 | " print(f\"\\nImage #{i+1}:\")\n", 1228 | " print(f\"True Label (CIFAR-10) = '{true_label}'\")\n", 1229 | " print(\"Predicted class Pretrained:\", model_pt.config.id2label[prediction_pt.item()])\n", 1230 | " print(\"Predicted class Ours:\", classes[prediction_ours])\n", 1231 | " print(\"-\" * 30)\n" 1232 | ] 1233 | }, 1234 | { 1235 | "cell_type": "markdown", 1236 | "metadata": { 1237 | "id": "UOrQmspBUJPr" 1238 | }, 1239 | "source": [ 1240 | "## Part 3: Self-supervised Learning via [MAEs](https://arxiv.org/abs/2111.06377)\n", 1241 | "\n", 1242 | "Above, you saw how we might train a vision encoder with labeled data using a classification loss function. But what happens if we do not have labels?\n", 1243 | "\n", 1244 | "There are many techniques that can be classified as \"self-supervised\" where we use the data itself to generate \"pre-text\" tasks that allow the model to learn representations from data without labels. One very common and popular method is known as Masked Autoencoders (MAE). The core idea is that a model will learn useful visual representations by having to reconstruct full images from partial (masked) views.\n", 1245 | "\n", 1246 | "
\n", 1247 | "\n", 1248 | "
\n", 1249 | "\n", 1250 | "In this section, we will delve into how this works and implement a simple example." 1251 | ] 1252 | }, 1253 | { 1254 | "cell_type": "markdown", 1255 | "metadata": { 1256 | "id": "gHF_iXI0yupE" 1257 | }, 1258 | "source": [ 1259 | "### Part 3.1: Implementing the mean-squared error loss\n", 1260 | "\n", 1261 | "The goal of the MAE is to reconstruct the original image from a partial, masked input. To measure how well the model is doing, we need a loss function. For image reconstruction, a common and effective choice is the **Mean Squared Error (MSE)**, which calculates the average squared difference between the predicted pixels and the true pixels.\n", 1262 | "\n", 1263 | "Your task is to implement this loss. You will compare the model's prediction with the original, unmasked image.\n", 1264 | "\n", 1265 | "$$ \\text{MSE} = \\frac{1}{N} \\sum_{i=1}^{N} (\\text{prediction}_i - \\text{target}_i)^2 $$" 1266 | ] 1267 | }, 1268 | { 1269 | "cell_type": "code", 1270 | "execution_count": null, 1271 | "metadata": { 1272 | "id": "2pDAlb7BzY-H" 1273 | }, 1274 | "outputs": [], 1275 | "source": [ 1276 | "#@title Hint 1 { display-mode: \"form\" }\n", 1277 | "# The task is to implement the Mean Squared Error (MSE) loss.\n", 1278 | "\n", 1279 | "# Calculate the squared difference between each pixel.\n", 1280 | "# Then, calculate the mean of all the squared differences.\n", 1281 | "# The .mean() operation on the entire tensor computes the average\n", 1282 | "# over all dimensions (e.g., N, C, H, W), resulting in a single scalar value.\n", 1283 | "\n", 1284 | "# Step 1: Solution\n", 1285 | "# squared_error = (predictions - targets) ** 2\n", 1286 | "# loss = torch.mean(squared_error)\n", 1287 | "\n", 1288 | "# Example\n", 1289 | "\n", 1290 | "# N = Batch Size, C = Channels, H = Height, W = Width\n", 1291 | "N, C, H, W = 4, 3, 32, 32\n", 1292 | "\n", 1293 | "# Create some dummy prediction and target tensors\n", 1294 | "predictions = torch.randn(N, C, H, W)\n", 1295 | "targets = torch.randn(N, C, H, W)\n", 1296 | "\n", 1297 | "# Calculate the squared difference between each element\n", 1298 | "squared_error = (predictions - targets) ** 2\n", 1299 | "\n", 1300 | "# Calculate the mean of all the squared differences\n", 1301 | "loss = torch.mean(squared_error)\n", 1302 | "\n", 1303 | "print(f\"Calculated MSE Loss: {loss.item()}\")\n", 1304 | "\n", 1305 | "# Verification with PyTorch's built-in MSELoss\n", 1306 | "mse_loss_fn = torch.nn.MSELoss()\n", 1307 | "torch_loss = mse_loss_fn(predictions, targets)\n", 1308 | "print(f\"PyTorch's MSE Loss: {torch_loss.item()}\")\n", 1309 | "\n", 1310 | "assert torch.allclose(loss, torch_loss), \"Your implementation does not match PyTorch's MSELoss.\"\n", 1311 | "print(\"\\nSuccess! Your implementation is correct.\")" 1312 | ] 1313 | }, 1314 | { 1315 | "cell_type": "code", 1316 | "execution_count": null, 1317 | "metadata": { 1318 | "id": "Ov2RERwmzOU4" 1319 | }, 1320 | "outputs": [], 1321 | "source": [ 1322 | "def reconstruction_loss(predictions, targets):\n", 1323 | " \"\"\"\n", 1324 | " Calculates the Mean Squared Error between predictions and targets.\n", 1325 | " This implementation does it manually to show the underlying operations.\n", 1326 | "\n", 1327 | " Args:\n", 1328 | " predictions (torch.Tensor): The output from the MAE decoder.\n", 1329 | " Shape: (N, C, H, W)\n", 1330 | " targets (torch.Tensor): The original, unmasked images.\n", 1331 | " Shape: (N, C, H, W)\n", 1332 | " Returns:\n", 1333 | " torch.Tensor: The reconstruction loss value (a scalar).\n", 1334 | " \"\"\"\n", 1335 | " # --- TODO 1: YOUR CODE HERE --- #\n", 1336 | " # Calculate the squared difference between each pixel\n", 1337 | "\n", 1338 | " # Calculate the mean of all the squared differences.\n", 1339 | " # The .mean() operation on the entire tensor computes the average\n", 1340 | " # over all dimensions (N, C, H, W), resulting in a single scalar value.\n", 1341 | "\n", 1342 | " return loss\n", 1343 | " # --- END YOUR CODE --- #" 1344 | ] 1345 | }, 1346 | { 1347 | "cell_type": "markdown", 1348 | "metadata": { 1349 | "id": "Xfs-7_vkz429" 1350 | }, 1351 | "source": [ 1352 | "### Part 3.2: Masking Images for the MAE\n", 1353 | "\n", 1354 | "This is the core mechanic of an MAE. We need a function that takes a batch of images and randomly masks a high percentage of their patches. In the original MAE paper, they use a 75% mask ratio.\n", 1355 | "\n", 1356 | "The process is:\n", 1357 | "1. Divide the image into a grid of non-overlapping patches.\n", 1358 | "2. Randomly select a certain percentage of these patches to mask.\n", 1359 | "3. Return the visible patches, the masked patches, and the indices that allow us to put them back together." 1360 | ] 1361 | }, 1362 | { 1363 | "cell_type": "code", 1364 | "execution_count": null, 1365 | "metadata": { 1366 | "id": "ikMHHaCRQPf2" 1367 | }, 1368 | "outputs": [], 1369 | "source": [ 1370 | "#@title Hint 1 { display-mode: \"form\" }\n", 1371 | "\n", 1372 | "# The task is to calculate the number of visible (unmasked) patches\n", 1373 | "# given a total number of patches and a mask ratio.\n", 1374 | "\n", 1375 | "# 'L' represents the total number of patches from an image.\n", 1376 | "# 'mask_ratio' is the fraction of patches to be removed or hidden.\n", 1377 | "# The result 'len_keep' is the number of patches that will remain visible.\n", 1378 | "\n", 1379 | "# Step 1: Solution\n", 1380 | "# len_keep = int(L * (1 - mask_ratio))\n", 1381 | "\n", 1382 | "\n", 1383 | "# Let's assume an image is split into a sequence of patches.\n", 1384 | "# For a 224x224 image with 16x16 patches, L = (224/16) * (224/16) = 14 * 14 = 196.\n", 1385 | "L = 196 # Total number of patches\n", 1386 | "mask_ratio = 0.75 # We want to mask 75% of the patches\n", 1387 | "\n", 1388 | "# --- Your implementation here ---\n", 1389 | "# Calculate the number of patches to keep (unmasked)\n", 1390 | "len_keep = int(L * (1 - mask_ratio))\n", 1391 | "# --- End of implementation ---\n", 1392 | "\n", 1393 | "print(f\"Total number of patches (L): {L}\")\n", 1394 | "print(f\"Mask ratio: {mask_ratio}\")\n", 1395 | "print(f\"Calculated number of visible patches (len_keep): {len_keep}\")\n", 1396 | "\n", 1397 | "# Verification\n", 1398 | "expected_len_keep = int(196 * (1 - 0.75))\n", 1399 | "assert len_keep == expected_len_keep, f\"Calculation is incorrect. Expected {expected_len_keep}.\"\n", 1400 | "print(\"\\nSuccess! Your implementation is correct.\")" 1401 | ] 1402 | }, 1403 | { 1404 | "cell_type": "code", 1405 | "execution_count": null, 1406 | "metadata": { 1407 | "id": "9CKj2dhyXQqV" 1408 | }, 1409 | "outputs": [], 1410 | "source": [ 1411 | "#@title Hint 2 { display-mode: \"form\" }\n", 1412 | "\n", 1413 | "# The task is to select a subset of items (visible patches) from a sequence\n", 1414 | "# and create a corresponding binary mask.\n", 1415 | "\n", 1416 | "# Step 1: Gather visible patches\n", 1417 | "# We use torch.gather to select patches from the input tensor 'x'\n", 1418 | "# based on the indices 'ids_keep'. The index tensor must be expanded\n", 1419 | "# to match the dimensions of 'x'.\n", 1420 | "# x_visible = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))\n", 1421 | "\n", 1422 | "# Step 2: Generate the binary mask\n", 1423 | "# Create a tensor of all ones, with the same batch size (N) and\n", 1424 | "# sequence length (L) as the original input.\n", 1425 | "# mask = torch.ones([N, L], device=x.device)\n", 1426 | "\n", 1427 | "# Step 3: Mark visible patches in the mask\n", 1428 | "# Use .scatter_ to place zeros at the locations of the visible patches.\n", 1429 | "# This modifies the mask in-place. The result is a mask where\n", 1430 | "# 0 indicates a visible patch and 1 indicates a masked patch.\n", 1431 | "# mask.scatter_(dim=1, index=ids_keep, value=0)\n", 1432 | "\n", 1433 | "# Example\n", 1434 | "\n", 1435 | "# N = Batch Size, L = Sequence Length (total patches), D = Dimension per patch\n", 1436 | "N, L, D = 2, 10, 8\n", 1437 | "# Number of patches to keep visible\n", 1438 | "len_keep = 4\n", 1439 | "\n", 1440 | "# Create a dummy input tensor (e.g., flattened image patches)\n", 1441 | "x = torch.randn(N, L, D)\n", 1442 | "\n", 1443 | "# Generate random indices for the patches to keep (visible patches)\n", 1444 | "# In a real scenario, these would be the result of a shuffling operation.\n", 1445 | "noise = torch.rand(N, L) # Noise to shuffle\n", 1446 | "ids_shuffle = torch.argsort(noise, dim=1) # Sort noise to get shuffled indices\n", 1447 | "ids_keep = ids_shuffle[:, :len_keep] # Keep the first 'len_keep' indices\n", 1448 | "\n", 1449 | "# Gather the visible patches using the selected indices\n", 1450 | "# We need to expand ids_keep to match the dimensions of x for gathering\n", 1451 | "index_for_gather = ids_keep.unsqueeze(-1).repeat(1, 1, D)\n", 1452 | "x_visible = torch.gather(x, dim=1, index=index_for_gather)\n", 1453 | "\n", 1454 | "# Generate the binary mask (0 for visible, 1 for masked)\n", 1455 | "mask = torch.ones([N, L], device=x.device)\n", 1456 | "# Use scatter_ to place 0s at the locations of the visible patches\n", 1457 | "mask.scatter_(dim=1, index=ids_keep, value=0)\n", 1458 | "\n", 1459 | "\n", 1460 | "print(\"--- Inputs ---\")\n", 1461 | "print(f\"Original sequence shape: {x.shape}\")\n", 1462 | "print(f\"Indices to keep (ids_keep):\\n{ids_keep}\")\n", 1463 | "\n", 1464 | "print(\"\\n--- Outputs ---\")\n", 1465 | "print(f\"Visible patches shape: {x_visible.shape}\")\n", 1466 | "print(f\"Binary mask (0=visible, 1=masked):\\n{mask}\")\n", 1467 | "\n", 1468 | "\n", 1469 | "# Verification\n", 1470 | "assert x_visible.shape == (N, len_keep, D), \"Shape of x_visible is incorrect.\"\n", 1471 | "assert mask.shape == (N, L), \"Shape of mask is incorrect.\"\n", 1472 | "# Verify that the number of visible patches (zeros) in the mask is correct\n", 1473 | "assert torch.sum(mask == 0).item() == N * len_keep, \"Mask does not have the correct number of visible patches.\"\n", 1474 | "print(\"\\nSuccess! Your implementation is correct.\")" 1475 | ] 1476 | }, 1477 | { 1478 | "cell_type": "code", 1479 | "execution_count": null, 1480 | "metadata": { 1481 | "id": "epioTw7cP4b1" 1482 | }, 1483 | "outputs": [], 1484 | "source": [ 1485 | "def random_masking(x, mask_ratio=0.75):\n", 1486 | " \"\"\"\n", 1487 | " Takes a batch of flattened patches and randomly masks them.\n", 1488 | " Args:\n", 1489 | " x (torch.Tensor): Input tensor of shape (N, L, D) where N is batch size,\n", 1490 | " L is sequence length (num patches), D is patch dimension.\n", 1491 | " mask_ratio (float): The ratio of patches to mask.\n", 1492 | " Returns:\n", 1493 | " tuple: A tuple containing:\n", 1494 | " - x_visible (torch.Tensor): The visible (unmasked) patches.\n", 1495 | " - mask (torch.Tensor): A binary mask indicating which patches were kept.\n", 1496 | " - ids_restore (torch.Tensor): Indices to restore the original order.\n", 1497 | " \"\"\"\n", 1498 | " N, L, D = x.shape\n", 1499 | "\n", 1500 | " # --- TODO 1: YOUR CODE HERE --- #\n", 1501 | " # Calculate the number of visible (unmasked) patches via the mask_ratio\n", 1502 | "\n", 1503 | "\n", 1504 | " # --- END YOUR CODE HERE --- #\n", 1505 | "\n", 1506 | " # Analogy: For each image in our batch, we shuffle a deck of cards (our patches).\n", 1507 | " # `torch.rand` creates random floats, `torch.argsort` gives us the indices that would\n", 1508 | " # sort them, which is a random permutation.\n", 1509 | "\n", 1510 | " noise = torch.rand(N, L, device=x.device) # A random value for each patch\n", 1511 | " ids_shuffle = torch.argsort(noise, dim=1) # The shuffled indices\n", 1512 | "\n", 1513 | " # This is needed to put the patches back in their original positions later\n", 1514 | " ids_restore = torch.argsort(ids_shuffle, dim=1)\n", 1515 | "\n", 1516 | " # Keep the first `len_keep` indices. These are our visible patches.\n", 1517 | " ids_keep = ids_shuffle[:, :len_keep]\n", 1518 | "\n", 1519 | " # --- TODO 2: YOUR CODE HERE --- #\n", 1520 | "\n", 1521 | " # Gather the visible patches using the selected indices\n", 1522 | "\n", 1523 | "\n", 1524 | " # Generate the binary mask for visualization and loss calculation\n", 1525 | " # 0 for visible patches, 1 for masked patches\n", 1526 | "\n", 1527 | "\n", 1528 | " # Put zeros for visible patches at the locations defined by ids_keep using\n", 1529 | " # the .scatter_ function\n", 1530 | "\n", 1531 | "\n", 1532 | " # --- END YOUR CODE HER --- #\n", 1533 | " return x_visible, mask, ids_restore\n", 1534 | "\n", 1535 | "# Helper function to visualize the masking process\n", 1536 | "def visualize_masking(patch_size=4, mask_ratio=0.75):\n", 1537 | " images, _ = next(iter(test_loader))\n", 1538 | " images = images.to(device)\n", 1539 | "\n", 1540 | " # For visualization, let's take just one image\n", 1541 | " img = images[0].unsqueeze(0)\n", 1542 | "\n", 1543 | " # 1. Patchify the image\n", 1544 | " patches = img.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)\n", 1545 | " patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(img.shape[0], -1, 3 * patch_size * patch_size)\n", 1546 | "\n", 1547 | " # 2. Perform masking\n", 1548 | " _, mask, _ = random_masking(patches, mask_ratio=mask_ratio)\n", 1549 | " mask = mask.detach().cpu().reshape(-1, img.shape[2]//patch_size, img.shape[3]//patch_size)\n", 1550 | "\n", 1551 | " # 3. Create the masked image for visualization\n", 1552 | " masked_img = img.clone().cpu()\n", 1553 | " for i in range(mask.shape[1]):\n", 1554 | " for j in range(mask.shape[2]):\n", 1555 | " if mask[0, i, j] == 1: # If the patch is masked\n", 1556 | " masked_img[0, :, i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size] = 0\n", 1557 | "\n", 1558 | " imshow(torchvision.utils.make_grid([images.cpu()[0], masked_img[0]]), title=\"Original vs. Masked\")\n", 1559 | "\n", 1560 | "print(\"Visualizing the masking process (75% of patches are blacked out):\")\n", 1561 | "visualize_masking(patch_size=4, mask_ratio=0.75)" 1562 | ] 1563 | }, 1564 | { 1565 | "cell_type": "markdown", 1566 | "metadata": { 1567 | "id": "_GTLhQYwq-mE" 1568 | }, 1569 | "source": [ 1570 | "### Part 3.3 The MAE Model Architecture\n", 1571 | "\n", 1572 | "Now we define the model itself. It consists of three main parts:\n", 1573 | "\n", 1574 | "1. **Patch + Positional Embedding:** This layer converts an image into a sequence of flattened patch embeddings. We also add positional embeddings to give the model information about where each patch is located in the original image.\n", 1575 | "2. **ViT Encoder:** A standard Vision Transformer encoder. It processes the sequence of *visible* patches and outputs their learned representations.\n", 1576 | "3. **ViT Decoder:** A lightweight Transformer decoder. It takes the encoded visible patches plus the mask tokens and reconstructs the full image.\n", 1577 | "\n", 1578 | "You have already seen the EncoderBlock and implemented parts of it. We will create a replica of that here for the DecoderBlock as well. Your next task will be to integrate these blocks within the main `MaskedAutoencoderViT` class." 1579 | ] 1580 | }, 1581 | { 1582 | "cell_type": "code", 1583 | "execution_count": null, 1584 | "metadata": { 1585 | "id": "iNrii0FuZJAN" 1586 | }, 1587 | "outputs": [], 1588 | "source": [ 1589 | "# This block is the same as the encoder block but we call it DecoderBlock here for readability\n", 1590 | "DecoderBlock = EncoderBlock" 1591 | ] 1592 | }, 1593 | { 1594 | "cell_type": "code", 1595 | "execution_count": null, 1596 | "metadata": { 1597 | "id": "yG5carEkcu9e" 1598 | }, 1599 | "outputs": [], 1600 | "source": [ 1601 | "#@title Hint 1 { display-mode: \"form\" }\n", 1602 | "\n", 1603 | "# The task is to append learnable mask tokens to a sequence of visible patch embeddings.\n", 1604 | "# This step reconstructs a full-length sequence to be fed into an encoder, like a Transformer,\n", 1605 | "# after the original input has had some of its patches removed.\n", 1606 | "\n", 1607 | "# Step 1: Create the mask tokens tensor.\n", 1608 | "# The mask_token is a learnable parameter, a single vector that will represent\n", 1609 | "# all masked-out patches. We repeat it to match the batch size and the number of\n", 1610 | "# masked patches.\n", 1611 | "# num_masked = self.num_patches - x_visible.shape[1]\n", 1612 | "# mask_tokens = self.mask_token.repeat(x_visible.shape[0], num_masked, 1)\n", 1613 | "\n", 1614 | "# Step 2: Concatenate visible patches and mask tokens.\n", 1615 | "# Combine the sequence of visible patch embeddings with the newly created\n", 1616 | "# mask tokens along the sequence dimension (dim=1).\n", 1617 | "# x_ = torch.cat([x_visible, mask_tokens], dim=1)\n", 1618 | "\n", 1619 | "# Example\n", 1620 | "\n", 1621 | "# --- Setup ---\n", 1622 | "# N = Batch Size, L = Total Sequence Length (total patches), D = Dimension per patch\n", 1623 | "N, L, D = 2, 16, 8\n", 1624 | "# Number of patches that are visible\n", 1625 | "len_keep = 4\n", 1626 | "\n", 1627 | "# Simulate the visible patch embeddings (output from a previous step)\n", 1628 | "x_visible = torch.randn(N, len_keep, D)\n", 1629 | "\n", 1630 | "# In a real model, mask_token is a learnable parameter.\n", 1631 | "# We'll simulate it here. It has a shape of (1, 1, D).\n", 1632 | "mask_token = torch.nn.Parameter(torch.zeros(1, 1, D))\n", 1633 | "# In a real scenario, this would be part of your model, e.g., self.mask_token.\n", 1634 | "\n", 1635 | "# --- Your implementation here ---\n", 1636 | "\n", 1637 | "# 1. Calculate the number of patches that were masked.\n", 1638 | "# The total number of patches is L. The number of visible patches is the\n", 1639 | "# length of the sequence dimension of x_visible.\n", 1640 | "num_masked = L - x_visible.shape[1]\n", 1641 | "\n", 1642 | "# Create the tensor of mask tokens by repeating the learnable mask_token.\n", 1643 | "# It should match the batch size (N) and the number of masked patches.\n", 1644 | "mask_tokens = mask_token.repeat(N, num_masked, 1)\n", 1645 | "\n", 1646 | "# 2. Concatenate the visible patch embeddings and the mask tokens.\n", 1647 | "x_ = torch.cat([x_visible, mask_tokens], dim=1)\n", 1648 | "\n", 1649 | "# --- End of implementation ---\n", 1650 | "\n", 1651 | "\n", 1652 | "print(\"--- Inputs ---\")\n", 1653 | "print(f\"Visible patches shape: {x_visible.shape}\")\n", 1654 | "print(f\"Learnable mask_token shape: {mask_token.shape}\")\n", 1655 | "print(f\"Number of total patches (L): {L}\")\n", 1656 | "print(f\"Number of visible patches: {len_keep}\")\n", 1657 | "print(f\"Number of masked patches: {num_masked}\")\n", 1658 | "\n", 1659 | "print(\"\\n--- Outputs ---\")\n", 1660 | "print(f\"Repeated mask_tokens shape: {mask_tokens.shape}\")\n", 1661 | "print(f\"Final concatenated sequence shape: {x_.shape}\")\n", 1662 | "\n", 1663 | "\n", 1664 | "# Verification\n", 1665 | "# The final sequence should have the original total length L.\n", 1666 | "assert x_.shape == (N, L, D), \"The final tensor shape is incorrect.\"\n", 1667 | "# The first part of the sequence should be the visible patches.\n", 1668 | "assert torch.all(x_[:, :len_keep, :] == x_visible), \"The visible patches were not concatenated correctly.\"\n", 1669 | "# The second part should be the mask tokens.\n", 1670 | "assert torch.all(x_[:, len_keep:, :] == mask_token), \"The mask tokens were not concatenated correctly.\"\n", 1671 | "\n", 1672 | "print(\"\\nSuccess! Your implementation is correct.\")" 1673 | ] 1674 | }, 1675 | { 1676 | "cell_type": "code", 1677 | "execution_count": null, 1678 | "metadata": { 1679 | "id": "QPpZamwgdN9s" 1680 | }, 1681 | "outputs": [], 1682 | "source": [ 1683 | "#@title Hint 2 { display-mode: \"form\" }\n", 1684 | "\n", 1685 | "# Step 1: Patchify the input images.\n", 1686 | "# Convert the batch of images into a sequence of flattened patches with positional embeddings.\n", 1687 | "# patches = self.patch_embed(imgs)\n", 1688 | "\n", 1689 | "# Step 2: Run the encoder.\n", 1690 | "# Randomly mask a large portion of the patches and feed only the\n", 1691 | "# visible patches through the encoder.\n", 1692 | "# x_visible, mask, ids_restore = self.forward_encoder(patches, self.mask_ratio)\n", 1693 | "\n", 1694 | "# Step 3: Run the decoder.\n", 1695 | "# Append learnable mask tokens to the encoded visible patches,\n", 1696 | "# restore the original sequence order, and run through the decoder\n", 1697 | "# to predict the pixel values for the full image.\n", 1698 | "# pred = self.forward_decoder(x_visible, ids_restore)\n", 1699 | "\n", 1700 | "# Step 4: Unpatchify the prediction.\n", 1701 | "# Reshape the sequence of predicted patches back into an image format.\n", 1702 | "# recon_imgs = self._unpatchify(pred)\n", 1703 | "\n", 1704 | "# Example\n", 1705 | "\n", 1706 | "\n", 1707 | "# --- Full MAE Model ---\n", 1708 | "class MaskedAutoencoderViTExample(nn.Module):\n", 1709 | " \"\"\"A simplified Masked Autoencoder with a Vision Transformer backbone.\"\"\"\n", 1710 | " def __init__(self, img_size=32, patch_size=4, in_chans=3,\n", 1711 | " embed_dim=64, encoder_depth=8, decoder_depth=2,\n", 1712 | " mlp_ratio=2.0, mask_ratio=0.75):\n", 1713 | " super().__init__()\n", 1714 | " self.patch_size = patch_size\n", 1715 | " self.mask_ratio = mask_ratio\n", 1716 | " self.num_patches = (img_size // patch_size) ** 2\n", 1717 | " patch_dim = in_chans * patch_size ** 2\n", 1718 | "\n", 1719 | " self.patch_embed = PatchAndEmbed(img_size, patch_size, in_chans, embed_dim)\n", 1720 | " self.encoder_blocks = nn.ModuleList([EncoderBlock(embed_dim, mlp_ratio) for _ in range(encoder_depth)])\n", 1721 | " self.encoder_norm = nn.LayerNorm(embed_dim)\n", 1722 | "\n", 1723 | " self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n", 1724 | " self.decoder_embed = nn.Linear(embed_dim, embed_dim)\n", 1725 | " self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))\n", 1726 | " self.decoder_blocks = nn.ModuleList([DecoderBlock(embed_dim, mlp_ratio) for _ in range(decoder_depth)])\n", 1727 | " self.decoder_norm = nn.LayerNorm(embed_dim)\n", 1728 | "\n", 1729 | " self.decoder_pred = nn.Linear(embed_dim, patch_dim)\n", 1730 | "\n", 1731 | " def _unpatchify(self, x):\n", 1732 | " p = self.patch_size\n", 1733 | " h = w = int(self.num_patches**0.5)\n", 1734 | " C = x.shape[-1] // (p * p)\n", 1735 | " x = x.reshape(x.shape[0], h, w, C, p, p)\n", 1736 | " x = x.permute(0, 3, 1, 4, 2, 5).reshape(x.shape[0], C, h * p, w * p)\n", 1737 | " return x\n", 1738 | "\n", 1739 | " def forward_encoder(self, x, mask_ratio):\n", 1740 | " x_visible, mask, ids_restore = random_masking(x, mask_ratio)\n", 1741 | " for blk in self.encoder_blocks:\n", 1742 | " x_visible = blk(x_visible)\n", 1743 | " x_visible = self.encoder_norm(x_visible)\n", 1744 | " return x_visible, mask, ids_restore\n", 1745 | "\n", 1746 | " def forward_decoder(self, x_visible, ids_restore):\n", 1747 | " x_visible = self.decoder_embed(x_visible)\n", 1748 | " num_masked = self.num_patches - x_visible.shape[1]\n", 1749 | " mask_tokens = self.mask_token.repeat(x_visible.shape[0], num_masked, 1)\n", 1750 | " x_ = torch.cat([x_visible, mask_tokens], dim=1)\n", 1751 | " x = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x_.shape[2]))\n", 1752 | " x = x + self.decoder_pos_embed\n", 1753 | " for blk in self.decoder_blocks:\n", 1754 | " x = blk(x)\n", 1755 | " x = self.decoder_norm(x)\n", 1756 | " x = self.decoder_pred(x)\n", 1757 | " return x\n", 1758 | "\n", 1759 | " def forward(self, imgs):\n", 1760 | " # 1. Patchify the input images.\n", 1761 | " patches = self.patch_embed(imgs)\n", 1762 | "\n", 1763 | " # 2. Run the encoder.\n", 1764 | " x_visible, mask, ids_restore = self.forward_encoder(patches, self.mask_ratio)\n", 1765 | "\n", 1766 | " # 3. Run the decoder.\n", 1767 | " pred_patches = self.forward_decoder(x_visible, ids_restore)\n", 1768 | "\n", 1769 | " # 4. Unpatchify the prediction to get the reconstructed image.\n", 1770 | " recon_imgs = self._unpatchify(pred_patches)\n", 1771 | "\n", 1772 | " # We also need the original images for the loss function, but here we just show the flow.\n", 1773 | " return recon_imgs, mask\n", 1774 | "\n", 1775 | "# --- Setup and Run ---\n", 1776 | "model = MaskedAutoencoderViTExample()\n", 1777 | "N, C, H, W = 2, 3, 32, 32\n", 1778 | "imgs = torch.randn(N, C, H, W)\n", 1779 | "\n", 1780 | "recon_imgs, mask = model(imgs)\n", 1781 | "\n", 1782 | "# --- Final Shapes ---\n", 1783 | "print(f\"Original images shape: {imgs.shape}\")\n", 1784 | "print(f\"Reconstructed images shape: {recon_imgs.shape}\")\n", 1785 | "print(f\"Mask shape: {mask.shape}\")\n", 1786 | "print(f\"\\nExample mask for first image (0=visible, 1=masked):\\n{mask[0].reshape(8, 8).int()}\")\n", 1787 | "\n", 1788 | "assert recon_imgs.shape == imgs.shape, \"Shape of reconstructed image is incorrect.\"\n", 1789 | "print(\"\\nSuccess! The forward pass flow is correct.\")" 1790 | ] 1791 | }, 1792 | { 1793 | "cell_type": "code", 1794 | "execution_count": null, 1795 | "metadata": { 1796 | "id": "RWe-zcT0ZdyR" 1797 | }, 1798 | "outputs": [], 1799 | "source": [ 1800 | "class MaskedAutoencoderViT(nn.Module):\n", 1801 | " \"\"\"A simplified Masked Autoencoder with a Vision Transformer backbone.\"\"\"\n", 1802 | " def __init__(self, img_size=32, patch_size=4, in_chans=3,\n", 1803 | " embed_dim=64, encoder_depth=8, decoder_depth=2,\n", 1804 | " mlp_ratio=2.0, mask_ratio=0.75):\n", 1805 | " super().__init__()\n", 1806 | " self.patch_size = patch_size\n", 1807 | " self.mask_ratio = mask_ratio\n", 1808 | " self.num_patches = (img_size // patch_size) ** 2\n", 1809 | " patch_dim = in_chans * patch_size ** 2\n", 1810 | "\n", 1811 | " # --- Patch + Positional Embeddings ---\n", 1812 | " self.patch_embed = PatchAndEmbed(img_size, patch_size, in_chans, embed_dim)\n", 1813 | "\n", 1814 | " # --- Encoder ---\n", 1815 | " self.encoder_blocks = nn.ModuleList([\n", 1816 | " EncoderBlock(embed_dim, mlp_ratio) for _ in range(encoder_depth)\n", 1817 | " ])\n", 1818 | " self.encoder_norm = nn.LayerNorm(embed_dim)\n", 1819 | "\n", 1820 | " # --- Decoder ---\n", 1821 | " self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n", 1822 | " self.decoder_embed = nn.Linear(embed_dim, embed_dim)\n", 1823 | " self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))\n", 1824 | "\n", 1825 | " self.decoder_blocks = nn.ModuleList([\n", 1826 | " DecoderBlock(embed_dim, mlp_ratio) for _ in range(decoder_depth)\n", 1827 | " ])\n", 1828 | " self.decoder_norm = nn.LayerNorm(embed_dim)\n", 1829 | "\n", 1830 | " # --- Final Projection Head ---\n", 1831 | " self.decoder_pred = nn.Linear(embed_dim, patch_dim)\n", 1832 | "\n", 1833 | " def _patchify(self, imgs):\n", 1834 | " p = self.patch_size\n", 1835 | " # (N, C, H, W) -> (N, L, C*P*P)\n", 1836 | " x = imgs.unfold(2, p, p).unfold(3, p, p)\n", 1837 | " x = x.permute(0, 2, 3, 1, 4, 5).reshape(imgs.shape[0], -1, 3 * p * p)\n", 1838 | " return x\n", 1839 | "\n", 1840 | " def _unpatchify(self, x):\n", 1841 | " p = self.patch_size\n", 1842 | " h = w = int(self.num_patches**0.5)\n", 1843 | " # (N, L, C*P*P) -> (N, C, H, W)\n", 1844 | " x = x.reshape(x.shape[0], h, w, 3, p, p)\n", 1845 | " x = x.permute(0, 3, 1, 4, 2, 5).reshape(x.shape[0], 3, h * p, w * p)\n", 1846 | " return x\n", 1847 | "\n", 1848 | " def forward_encoder(self, x, mask_ratio):\n", 1849 | "\n", 1850 | " # Masking\n", 1851 | " x_visible, mask, ids_restore = random_masking(x, mask_ratio)\n", 1852 | "\n", 1853 | " # Run through encoder blocks\n", 1854 | " for blk in self.encoder_blocks:\n", 1855 | " x_visible = blk(x_visible)\n", 1856 | " x_visible = self.encoder_norm(x_visible)\n", 1857 | "\n", 1858 | " return x_visible, mask, ids_restore\n", 1859 | "\n", 1860 | " def forward_decoder(self, x_visible, ids_restore):\n", 1861 | " x_visible = self.decoder_embed(x_visible)\n", 1862 | "\n", 1863 | " # --- TODO 1: YOUR CODE HERE --- #\n", 1864 | " # 1. Create the mask tokens tensor by repeating the learnable `mask_token`.\n", 1865 | " # The number of mask tokens needed is `num_patches - num_visible_patches`.\n", 1866 | "\n", 1867 | "\n", 1868 | "\n", 1869 | " # 2. Concatenate the visible patch embeddings and the mask tokens together.\n", 1870 | "\n", 1871 | "\n", 1872 | " # --- END YOUR CODE --- #\n", 1873 | "\n", 1874 | " # 3. Use `ids_restore` to unshuffle the sequence back to its original patch order.\n", 1875 | " x = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x_.shape[2]))\n", 1876 | "\n", 1877 | " # Add decoder positional embeddings\n", 1878 | " x = x + self.decoder_pos_embed\n", 1879 | "\n", 1880 | " # Run through decoder blocks\n", 1881 | " for blk in self.decoder_blocks:\n", 1882 | " x = blk(x)\n", 1883 | " x = self.decoder_norm(x)\n", 1884 | "\n", 1885 | " # Final prediction head\n", 1886 | " x = self.decoder_pred(x)\n", 1887 | " return x\n", 1888 | "\n", 1889 | " def forward(self, imgs):\n", 1890 | " # --- TODO 2: YOUR CODE HERE --- #\n", 1891 | " # 1. Patchify the input images.\n", 1892 | "\n", 1893 | "\n", 1894 | " # 2. Run the encoder.\n", 1895 | "\n", 1896 | "\n", 1897 | " # 3. Run the decoder.\n", 1898 | "\n", 1899 | "\n", 1900 | " # 4. Unpatchify the prediction to get the reconstructed image.\n", 1901 | "\n", 1902 | "\n", 1903 | " # --- END YOUR CODE --- #\n", 1904 | "\n", 1905 | " # We also need the original images for the loss function.\n", 1906 | " return recon_imgs, imgs\n" 1907 | ] 1908 | }, 1909 | { 1910 | "cell_type": "markdown", 1911 | "metadata": { 1912 | "id": "z7WHA6KwdpTc" 1913 | }, 1914 | "source": [ 1915 | "## Part 3.4: Visualize the Reconstructed Output (Before Training)\n", 1916 | "\n", 1917 | "Let's see what our randomly initialized model produces. We'll pass a batch of images through it and visualize the output. Since the model hasn't learned anything, the reconstruction will look like random noise. This provides a good baseline to see how much the model improves with training." 1918 | ] 1919 | }, 1920 | { 1921 | "cell_type": "code", 1922 | "execution_count": null, 1923 | "metadata": { 1924 | "id": "XFgULL3Zd4OL" 1925 | }, 1926 | "outputs": [], 1927 | "source": [ 1928 | "def show_reconstruction(model, test_loader, device):\n", 1929 | " model.eval() # Set model to evaluation mode\n", 1930 | " images, _ = next(iter(test_loader))\n", 1931 | " images = images.to(device)\n", 1932 | "\n", 1933 | " with torch.no_grad():\n", 1934 | " recon_imgs, orig_imgs = model(images)\n", 1935 | "\n", 1936 | " # Create masked image for visualization\n", 1937 | " # This requires running parts of the model again to get the mask\n", 1938 | " patches = model._patchify(images)\n", 1939 | " _, mask, _ = random_masking(patches, model.mask_ratio)\n", 1940 | " mask = mask.detach().reshape(-1, model.num_patches)\n", 1941 | "\n", 1942 | " # Recreate the masked images\n", 1943 | " masked_patches = patches.clone()\n", 1944 | " for i in range(masked_patches.shape[0]):\n", 1945 | " masked_patches[i, mask[i]==1, :] = 0 # Black out masked patches\n", 1946 | " masked_imgs = model._unpatchify(masked_patches)\n", 1947 | "\n", 1948 | " # Move to CPU and create grids for display\n", 1949 | " orig_grid = torchvision.utils.make_grid(orig_imgs.cpu())\n", 1950 | " masked_grid = torchvision.utils.make_grid(masked_imgs.cpu())\n", 1951 | " recon_grid = torchvision.utils.make_grid(recon_imgs.cpu())\n", 1952 | "\n", 1953 | " plt.figure(figsize=(8,8))\n", 1954 | " imshow(orig_grid, title='Original Images')\n", 1955 | " plt.figure(figsize=(8,8))\n", 1956 | " imshow(masked_grid, title='Masked Images')\n", 1957 | " plt.figure(figsize=(8,8))\n", 1958 | " imshow(recon_grid, title='Reconstructed Images')\n", 1959 | " model.train() # Set model back to training mode\n", 1960 | "\n", 1961 | "# Instantiate the model\n", 1962 | "model_mae = MaskedAutoencoderViT(\n", 1963 | " img_size=32,\n", 1964 | " patch_size=4,\n", 1965 | " embed_dim=64,\n", 1966 | " encoder_depth=4,\n", 1967 | " decoder_depth=1,\n", 1968 | " mlp_ratio=2.0,\n", 1969 | " mask_ratio=0.75\n", 1970 | ").to(device)\n", 1971 | "\n", 1972 | "print(\"Model architecture:\", model_mae)\n", 1973 | "print(\"\\nVisualizing initial random reconstruction:\")\n", 1974 | "show_reconstruction(model_mae, test_loader, device)" 1975 | ] 1976 | }, 1977 | { 1978 | "cell_type": "markdown", 1979 | "metadata": { 1980 | "id": "9D7QfbzveJYl" 1981 | }, 1982 | "source": [ 1983 | "### Part 3.5: Running a training loop and visualizing reconstructions!\n", 1984 | "\n", 1985 | "Now that we have all the pieces, we can train a simple ViT based MAE. After training, the model learns to reconstruct images better and better. At convergence, the encoder will have learned to represent images in a latent embedding vector that is useful for image reconstruction via a shallow decoder. By doing this, this embedding can then be used for other tasks just like from supervised training." 1986 | ] 1987 | }, 1988 | { 1989 | "cell_type": "code", 1990 | "execution_count": null, 1991 | "metadata": { 1992 | "id": "g0TrYaZsetpL" 1993 | }, 1994 | "outputs": [], 1995 | "source": [ 1996 | "# Training Configuration\n", 1997 | "epochs = 4 # Increase this for better results\n", 1998 | "lr = 1.5e-4\n", 1999 | "wd=0.05\n", 2000 | "\n", 2001 | "optimizer = torch.optim.AdamW(model_mae.parameters(), lr=lr, weight_decay=0.05)\n", 2002 | "\n", 2003 | "print(\"Starting training...\")\n", 2004 | "model_mae.train()\n", 2005 | "\n", 2006 | "for epoch in range(epochs):\n", 2007 | " total_loss = 0\n", 2008 | " pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{epochs}\")\n", 2009 | " for imgs, _ in pbar:\n", 2010 | " imgs = imgs.to(device)\n", 2011 | "\n", 2012 | " # Forward pass\n", 2013 | " recon_imgs, orig_imgs = model_mae(imgs)\n", 2014 | " loss = reconstruction_loss(recon_imgs, orig_imgs)\n", 2015 | "\n", 2016 | " # Backward pass and optimization\n", 2017 | " optimizer.zero_grad()\n", 2018 | " loss.backward()\n", 2019 | " optimizer.step()\n", 2020 | "\n", 2021 | " total_loss += loss.item()\n", 2022 | " pbar.set_postfix({'loss': loss.item()})\n", 2023 | "\n", 2024 | " avg_loss = total_loss / len(train_loader)\n", 2025 | " print(f\"Epoch {epoch+1}/{epochs} - Average Loss: {avg_loss:.4f}\")\n", 2026 | "\n", 2027 | " # Visualize progress every 5 epochs\n", 2028 | " if (epoch + 1) % 1 == 0:\n", 2029 | " print(f\"\\n--- Visualizing reconstruction after epoch {epoch+1} ---\")\n", 2030 | " show_reconstruction(model_mae, test_loader, device)\n", 2031 | " print(\"--------------------------------------------------\\n\")\n", 2032 | "\n", 2033 | "print(\"Training finished!\")" 2034 | ] 2035 | }, 2036 | { 2037 | "cell_type": "markdown", 2038 | "source": [ 2039 | "### Part 3.6 Using a Pre-trained MAE representation for tasks!\n", 2040 | "\n", 2041 | "Previously we trained our CIFAR-10 ViT classifier from scratch. In 1 epoch we achieved 35% accuracy. What if we initialized our vision encoder from the encoder within the MAE we just trained? We can see that this encoder must have learned about the general structure of images and *encoded* that in it's 64 dimensional embedding. Now we'll see how well our classifier does if we train it for 1 epoch but with the encoder initialized with this pre-trained encoder.\n", 2042 | "\n", 2043 | "We will first create a class to wrap a pretrained patch embedding and encoder with a classifier head. We then initialize the patch_embedding, and encoder blocks to those that we just trained with the MAE." 2044 | ], 2045 | "metadata": { 2046 | "id": "He_J_BG2irSf" 2047 | } 2048 | }, 2049 | { 2050 | "cell_type": "code", 2051 | "source": [ 2052 | "class ClassifyWithMAE(nn.Module):\n", 2053 | " def __init__(self, patch_layer, encoder, embed_dim=64, num_classes = 10):\n", 2054 | " super().__init__()\n", 2055 | " self.patch_and_embed = patch_layer\n", 2056 | " self.encoder = encoder\n", 2057 | "\n", 2058 | " # Classification Head\n", 2059 | " self.norm = nn.LayerNorm(embed_dim)\n", 2060 | " self.head = nn.Linear(embed_dim, num_classes)\n", 2061 | "\n", 2062 | " def forward(self, x):\n", 2063 | " x = self.patch_and_embed(x)\n", 2064 | " for blk in self.encoder:\n", 2065 | " x = blk(x)\n", 2066 | " # After the final block, the sequence is normalised\n", 2067 | " x = self.norm(x)\n", 2068 | "\n", 2069 | " x = x.mean(dim=1)\n", 2070 | " logits = self.head(x)\n", 2071 | "\n", 2072 | " return logits\n", 2073 | "\n", 2074 | "model_mae_cls = ClassifyWithMAE(model_mae.patch_embed, model_mae.encoder_blocks, embed_dim=64, num_classes=10)\n", 2075 | "model_mae_cls.to(device)\n" 2076 | ], 2077 | "metadata": { 2078 | "id": "9QuTqtgofRYs" 2079 | }, 2080 | "execution_count": null, 2081 | "outputs": [] 2082 | }, 2083 | { 2084 | "cell_type": "code", 2085 | "source": [ 2086 | "criterion = nn.CrossEntropyLoss()\n", 2087 | "optimizer = torch.optim.AdamW(model_mae_cls.parameters(), lr=1.5e-4)\n", 2088 | "epochs=1\n", 2089 | "print(\"\\nStarting training for one epoch...\")\n", 2090 | "epoch = 0\n", 2091 | "for epoch in range(epochs):\n", 2092 | " avg_loss, accuracy = train_one_epoch(model_mae_cls, train_loader, criterion, optimizer, device, epoch)\n", 2093 | " print(f\"\\n--- End of Epoch ---\")\n", 2094 | " print(f\"Average Training Loss: {avg_loss:.4f}\")\n", 2095 | " print(f\"Training Accuracy: {accuracy:.2f}%\")" 2096 | ], 2097 | "metadata": { 2098 | "id": "xUSpuFy4g_l1" 2099 | }, 2100 | "execution_count": null, 2101 | "outputs": [] 2102 | }, 2103 | { 2104 | "cell_type": "markdown", 2105 | "source": [ 2106 | "In the supevised training of the same ViT architecture from scratch that we did earlier: [Supervised Training](#scrollTo=E3MxL7AWaVtK&line=52&uniqifier=1), we saw that in 1 epoch we achieved a training accuracy of approximately 34%. Now by initializing the encoder from the MAE we get a 4% boost! This shows that pretraining (even with a self-supervised task unrelated to classification) can improve the efficiency of learning during task adaptation!" 2107 | ], 2108 | "metadata": { 2109 | "id": "Z1mBR-O9oRXv" 2110 | } 2111 | }, 2112 | { 2113 | "cell_type": "markdown", 2114 | "metadata": { 2115 | "id": "GJiyM1KvcBO7" 2116 | }, 2117 | "source": [ 2118 | "### Part 4 (Take-home task): Contrastive Language-Image Pre-training (CLIP)\n", 2119 | "\n", 2120 | "CLIP aims to learn visual concepts through large-scale image-text pairs, widely available on the web. It operates by training a text encoder and image encoder to match captions with corresponding images.\n", 2121 | "\n", 2122 | "Specifically, given a batch of image-text pairs, CLIP computes the dense cosine similarity matrix between all possible (image, text) candidates within this batch. The core idea is to **maximize the similarity between the correct pairs** (shown in blue in the figure below) and **minimize the similarity for incorrect pairs** (shown in grey in the image). To do it, they optimize a symmetric cross-entropy loss over these similarity scores.\n", 2123 | "\n", 2124 | "
\n", 2125 | "\n", 2126 | "
\n", 2127 | "\n", 2128 | "Note that the Image Encoder and Text Encoder work together to project both images and text into a single, **multi-modal embedding space** where they can be compared.\n", 2129 | "\n", 2130 | "CLIP's ability to connect language and vision unlocks a wide range of powerful applications:\n", 2131 | "\n", 2132 | "- **Image and Text Similarity Search**: Find images that best match a text query or find text that best describes an image.\n", 2133 | "\n", 2134 | "- **Vision-Language Modelling**: The trained CLIP Vision Encoder can be used to embed images and train LLMs on image-text data for applications such as visual question answering and image captioning (more details on this in tomorrow's lecture on VLMs).\n", 2135 | "\n", 2136 | "To get hands-on with CLIP, you can work through the exercise from last year's [tutorial](https://colab.sandbox.google.com/github/eemlcommunity/PracticalSessions2024/blob/main/3_vision_language_models/VLM_tutorial.ipynb#scrollTo=Lugc4X-2Ls9g).\n", 2137 | "\n" 2138 | ] 2139 | }, 2140 | { 2141 | "cell_type": "markdown", 2142 | "metadata": { 2143 | "id": "T02yjuJ1apHb" 2144 | }, 2145 | "source": [ 2146 | "### Part IV: Conclusion\n", 2147 | "\n", 2148 | "Congratulations on completing this deep dive into training Vision Transformer encoders! 🥳 We covered:\n", 2149 | "\n", 2150 | "- **Fully Supervised Learning**: We began with the classic approach, teaching models to map images to explicit, human-provided labels.\n", 2151 | "\n", 2152 | "- **Self-Supervised Learning (SSL)**: We then explored the clever world of Masked Autoencoders (MAE), where models learn rich features by solving the \"puzzle\" of reconstructing parts of an image.\n", 2153 | "\n", 2154 | "- **Contrastive Learning**: Finally, we saw how models like CLIP can align vision and language in a shared semantic space by learning from massive, noisy datasets of image-text pairs.\n", 2155 | "\n", 2156 | "By mastering these methodologies, you have grasped the foundational principles behind most state-of-the-art vision backbones.\n", 2157 | "\n", 2158 | "In practice, the lines between these approaches are blurring. State-of-the-art research frequently combines them using hybrid loss functions to get the best of all worlds. For instance, a model might combine:\n", 2159 | "\n", 2160 | "- A **contrastive loss** (like CLIP's) to learn high-level semantic understanding.\n", 2161 | "\n", 2162 | "- A **reconstruction loss** (like MAE's) to ensure it also captures fine-grained, pixel-level details.\n", 2163 | "\n", 2164 | "The CoCa paper, which simultaneously performs contrastive and captioning losses, is a prime example.\n", 2165 | "\n", 2166 | "\n" 2167 | ] 2168 | }, 2169 | { 2170 | "cell_type": "markdown", 2171 | "metadata": { 2172 | "id": "OlTBiuybNvLz" 2173 | }, 2174 | "source": [ 2175 | "### Recommended Reading\n", 2176 | "\n", 2177 | "\n", 2178 | "To continue your journey, we highly recommend exploring the original papers that introduced these groundbreaking ideas and related works.\n", 2179 | "\n", 2180 | "- ViT: Dosovitskiy et al. (2020). [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.](https://arxiv.org/abs/2010.11929)\n", 2181 | "\n", 2182 | "- MAE: He et al. (2021). [Masked Autoencoders Are Scalable Vision Learners.](https://arxiv.org/abs/2111.06377)\n", 2183 | "\n", 2184 | "- DINO: Caron et al. (2021). [Emerging Properties in Self-Supervised Vision Transformers.](https://arxiv.org/abs/2104.14294)\n", 2185 | "\n", 2186 | "- CLIP: Radford et al. (2021). [Learning Transferable Visual Models From Natural Language Supervision.](https://arxiv.org/abs/2103.00020)\n", 2187 | "\n", 2188 | "- CoCa: Yu et al. (2022). [CoCa: Contrastive Captioners are Image-Text Foundation Models.](https://arxiv.org/abs/2205.01917)\n", 2189 | "\n", 2190 | "- Swin Transformer: Liu et al. (2021). [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows.](https://arxiv.org/abs/2103.14030)\n", 2191 | "\n", 2192 | "- BEiT: Bao et al. (2021). [BEiT: BERT Pre-Training of Image Transformers.](https://arxiv.org/abs/2106.08254)\n", 2193 | "\n", 2194 | "- SimCLR: Chen et al. (2020). [A Simple Framework for Contrastive Learning of Visual Representations.](https://arxiv.org/abs/2002.05709)\n", 2195 | "\n", 2196 | "- MoCo: He et al. (2019). [Momentum Contrast for Unsupervised Visual Representation Learning.](https://arxiv.org/abs/1911.05722)" 2197 | ] 2198 | } 2199 | ], 2200 | "metadata": { 2201 | "colab": { 2202 | "provenance": [] 2203 | }, 2204 | "kernelspec": { 2205 | "display_name": "Python 3", 2206 | "name": "python3" 2207 | }, 2208 | "language_info": { 2209 | "name": "python" 2210 | } 2211 | }, 2212 | "nbformat": 4, 2213 | "nbformat_minor": 0 2214 | } -------------------------------------------------------------------------------- /2_mechanistic_interpretability/README.md: -------------------------------------------------------------------------------- 1 | # [[EEML2025](https://www.eeml.eu)] Tutorial 2: Mechanistic Interpretability 2 | 3 | **Authors:** Arthur Conmy and Federico Barbero 4 | 5 | 6 | --- 7 | 8 | The goal of mechanistic interpretability is to take a trained model and reverse engineer the algorithms the model learned during training from its weights. It is a fact about the world today that we have computer programs that can essentially speak English at a human level (GPT-3 was announced in 2020, PaLM in 2022, many more LLMs since then), and yet we have no idea how they work nor how to write one ourselves. The goal of mechanistic interpretability is to try to fix this! 9 | 10 | ### Outline: 11 | 12 | - 1️⃣ TransformerLens: Introduction 13 | - 2️⃣ TransformerLens: Hooks 14 | - 3️⃣ Bonus: Replicating Golden Gate Claude! 15 | 16 | 17 | ### Notebooks 18 | 19 | Tutorial: [![Open In 20 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/2_mechanistic_interpretability/mechanistic_interpretability_tutorial.ipynb) 21 | 22 | --- 23 | -------------------------------------------------------------------------------- /3_rl/README.md: -------------------------------------------------------------------------------- 1 | # [[EEML2025](https://www.eeml.eu)] Tutorial 3: Reinforcement Learning 2 | 3 | **Authors:** Daniele Calandriello and Miruna Pîslar 4 | 5 | 6 | --- 7 | 8 | Reinforcement Learning (RL) is about learning to make good decisions through trial and error. An agent (our learner) interacts with an environment (the world or game it's in) by taking actions. After each action, the environment provides a reward (a score) and a new state (the new observation of the world). 9 | 10 | ### Outline: 11 | 12 | - Setup & Installation 13 | - Intro to Reinforcement Learning: an overview of the main RL concepts. 14 | - Practical 1: build and train an agent from scratch using the REINFORCE algorithm. 15 | - Practical 2 (Advanced): explore more powerful algorithms (A2C and PPO) by introducing concepts like value functions, advantage estimation, and entropy. 16 | - Practical 3: learn about RL*F (Reinforcement Learning from Feedback) and use it to fine-tune GPT-2 to generate more positive text. 17 | 18 | 19 | ### Notebooks 20 | 21 | Tutorial: [![Open In 22 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/3_rl/reinforcement_learning_tutorial.ipynb) 23 | 24 | Tutorial with solutions: [![Open In 25 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/3_rl/reinforcement_learning_tutorial_solutions.ipynb) 26 | 27 | 28 | --- 29 | -------------------------------------------------------------------------------- /4_flow_matching_drug_discovery/README.md: -------------------------------------------------------------------------------- 1 | # [[EEML2025](https://www.eeml.eu)] Tutorial 4: Flow Matching & Drug Discovery 2 | 3 | **Authors:** Oscar Davis and Katarina Petrović 4 | 5 | 6 | --- 7 | 8 | `flow_matching.ipynb` 9 | 10 | Welcome to the flow matching notebook, accompanying Joey Bose's talk! Here, we shall present a quick but thorough tour of what flow matching is and how it can be (rather) easily ported to specific geometries (Riemannian manifolds). In the first part, we shall see how to perform usual flow matching, with various samplers, and various couplings allowing for faster training. In the second part, we shall show how to port all of this to the mentioned geometries. 11 | 12 | 13 | `drug_discovery.ipynb` 14 | In this part of tutorial we will focus on applying our knowledge to generate cell state trajectories in single-cell data 🧬🧫 These trajectories are cool because they reveal tissue mechanics and disease progression, making them crucial for drug discovery and medical research!🏥 15 | 16 | Single-cell technology has changed the way we study biology. Instead of averaging signals from millions of cells, we can now measure thousands of cells one by one. 17 | 18 | 19 | ### Outline: 20 | 21 | 22 | `flow_matching.ipynb` 23 | - Setup 24 | - Flow matching in ℝ^𝕟 25 | - Flow matching on Riemannian manifolds 26 | 27 | `drug_discovery.ipynb` 28 | - Learning Cell State Transition Paths in Single-Cell Data 29 | - Metric Flow Matching 30 | 31 | 32 | 33 | ### Notebooks 34 | 35 | Tutorial part 1 - flow matching: [![Open In 36 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/4_flow_matching_drug_discovery/flow_matching.ipynb) 37 | 38 | Tutorial part 2 - drug discovery: [![Open In 39 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/4_flow_matching_drug_discovery/drug_discovery.ipynb) 40 | 41 | Tutorial part 1 with solutions - flow matching: [![Open In 42 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/4_flow_matching_drug_discovery/flow_matching_solutions.ipynb) 43 | 44 | Tutorial part 2 with solutions - drug discovery: [![Open In 45 | Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eemlcommunity/PracticalSessions2025/blob/main/4_flow_matching_drug_discovery/drug_discovery_solutions.ipynb) 46 | 47 | --- 48 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PracticalSessions2025 2 | 3 | Repository for tutorial sessions at EEML2025 in Sarajevo. 4 | 5 | Designed for education purposes. Please do not distribute without permission. Write to contact@eeml.eu if you have any questions. 6 | 7 | You are welcome to reuse this material in other courses or schools, but please reach out to contact@eeml.eu if you plan to do so. We would appreciate it if you could acknowledge that the materials come from EEML2025 and give credits to the authors. Also please keep a link in your materials to the original repo, in case updates occur. 8 | 9 | MIT License 10 | 11 | Copyright (c) 2025 EEML 12 | 13 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. --------------------------------------------------------------------------------