{"cells": [{"cell_type": "markdown", "id": "b15ebccc", "metadata": {"papermill": {"duration": 0.021786, "end_time": "2025-04-03T19:34:48.860019", "exception": false, "start_time": "2025-04-03T19:34:48.838233", "status": "completed"}, "tags": []}, "source": ["\n", "# Tutorial 12: Meta-Learning - Learning to Learn\n", "\n", "* **Author:** Phillip Lippe\n", "* **License:** CC BY-SA\n", "* **Generated:** 2025-04-03T19:34:42.067657\n", "\n", "In this tutorial, we will discuss algorithms that learn models which can quickly adapt to new classes and/or tasks with few samples.\n", "This area of machine learning is called _Meta-Learning_ aiming at \"learning to learn\".\n", "Learning from very few examples is a natural task for humans. In contrast to current deep learning models, we need to see only a few examples of a police car or firetruck to recognize them in daily traffic.\n", "This is crucial ability since in real-world application, it is rarely the case that the data stays static and does not change over time.\n", "For example, an object detection system for mobile phones trained on data from 2000 will have troubles detecting today's common mobile phones, and thus, needs to adapt to new data without excessive label effort.\n", "The optimization techniques we have discussed so far struggle with this because they only aim at obtaining a good performance on a test set that had similar data.\n", "However, what if the test set has classes that we do not have in the training set?\n", "Or what if we want to test the model on a completely different task?\n", "We will discuss and implement three common Meta-Learning algorithms for such situations.\n", "This notebook is part of a lecture series on Deep Learning at the University of Amsterdam.\n", "The full list of tutorials can be found at https://uvadlc-notebooks.rtfd.io.\n", "\n", "\n", "---\n", "Open in [{height=\"20px\" width=\"117px\"}](https://colab.research.google.com/github/PytorchLightning/lightning-tutorials/blob/publication/.notebooks/course_UvA-DL/12-meta-learning.ipynb)\n", "\n", "Give us a \u2b50 [on Github](https://www.github.com/Lightning-AI/lightning/)\n", "| Check out [the documentation](https://lightning.ai/docs/)\n", "| Join us [on Discord](https://discord.com/invite/tfXFetEZxv)"]}, {"cell_type": "markdown", "id": "3c7e16ba", "metadata": {"papermill": {"duration": 0.015974, "end_time": "2025-04-03T19:34:48.897937", "exception": false, "start_time": "2025-04-03T19:34:48.881963", "status": "completed"}, "tags": []}, "source": ["## Setup\n", "This notebook requires some packages besides pytorch-lightning."]}, {"cell_type": "code", "execution_count": 1, "id": "8a2d3960", "metadata": {"colab": {}, "colab_type": "code", "execution": {"iopub.execute_input": "2025-04-03T19:34:48.924585Z", "iopub.status.busy": "2025-04-03T19:34:48.923867Z", "iopub.status.idle": "2025-04-03T19:34:50.125860Z", "shell.execute_reply": "2025-04-03T19:34:50.124566Z"}, "id": "LfrJLKPFyhsK", "lines_to_next_cell": 0, "papermill": {"duration": 1.218057, "end_time": "2025-04-03T19:34:50.128446", "exception": false, "start_time": "2025-04-03T19:34:48.910389", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable.It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\r\n", "\u001b[0m"]}, {"name": "stdout", "output_type": "stream", "text": ["\r\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\r\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython -m pip install --upgrade pip\u001b[0m\r\n"]}], "source": ["! pip install --quiet \"torchmetrics >=1.0,<1.8\" \"tensorboard\" \"scipy\" \"torch >=1.8.1,<2.7\" \"pytorch-lightning >=2.0,<2.6\" \"matplotlib\" \"numpy <3.0\" \"seaborn\" \"torchvision\""]}, {"cell_type": "markdown", "id": "b0ba385a", "metadata": {"papermill": {"duration": 0.020125, "end_time": "2025-04-03T19:34:50.170026", "exception": false, "start_time": "2025-04-03T19:34:50.149901", "status": "completed"}, "tags": []}, "source": ["
\n", "Meta-Learning offers solutions to these situations, and we will discuss three popular algorithms: __Prototypical Networks__ ([Snell et al., 2017](https://arxiv.org/abs/1703.05175)), __Model-Agnostic Meta-Learning / MAML__ ([Finn et al., 2017](http://proceedings.mlr.press/v70/finn17a.html)), and __Proto-MAML__ ([Triantafillou et al., 2020](https://openreview.net/pdf?id=rkgAGAVKPr)).\n", "We will focus on the task of few-shot classification where the training and test set have distinct sets of classes.\n", "For instance, we would train the model on the binary classifications of cats-birds and flowers-bikes, but during test time, the model would need to learn from 4 examples each the difference between dogs and otters, two classes we have not seen during training (Figure credit - [Lilian Weng](https://lilianweng.github.io/lil-log/2018/11/30/meta-learning.html)).\n", "\n", "
\n", "\n", "A different setup, which is very common in Reinforcement Learning and recently Natural Language Processing, is to aim at few-shot learning of a completely new task.\n", "For example, an robot agent that learned to run, jump and pick up boxes, should quickly adapt to collecting and stacking boxes.\n", "In NLP, we can think of a model which was trained sentiment classification, hatespeech detection and sarcasm classification, to adapt to classifying the emotion of a text.\n", "All methods we will discuss in this notebook can be easily applied to these settings since we only use a different definition of a 'task'.\n", "For few-shot classification, we consider a task to distinguish between $M$ novel classes.\n", "Here, we would not only have novel classes, but also a completely different dataset.\n", "\n", "First of all, let's start with importing our standard libraries. We will again be using PyTorch Lightning."]}, {"cell_type": "code", "execution_count": 2, "id": "37cca226", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:34:50.200605Z", "iopub.status.busy": "2025-04-03T19:34:50.200227Z", "iopub.status.idle": "2025-04-03T19:34:54.243813Z", "shell.execute_reply": "2025-04-03T19:34:54.242929Z"}, "papermill": {"duration": 4.059088, "end_time": "2025-04-03T19:34:54.245115", "exception": false, "start_time": "2025-04-03T19:34:50.186027", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["Seed set to 42\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Device: cuda:0\n"]}, {"data": {"text/plain": [""]}, "metadata": {}, "output_type": "display_data"}], "source": ["import json\n", "import os\n", "import random\n", "import urllib.request\n", "from collections import defaultdict\n", "from copy import deepcopy\n", "from statistics import mean, stdev\n", "from urllib.error import HTTPError\n", "\n", "import matplotlib\n", "import matplotlib.pyplot as plt\n", "import matplotlib_inline.backend_inline\n", "import numpy as np\n", "import pytorch_lightning as pl\n", "import seaborn as sns\n", "import torch\n", "import torch.nn.functional as F\n", "import torch.optim as optim\n", "import torch.utils.data as data\n", "import torchvision\n", "from PIL import Image\n", "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n", "from torchvision import transforms\n", "from torchvision.datasets import CIFAR100, SVHN\n", "from tqdm.auto import tqdm\n", "\n", "plt.set_cmap(\"cividis\")\n", "%matplotlib inline\n", "matplotlib_inline.backend_inline.set_matplotlib_formats(\"svg\", \"pdf\") # For export\n", "matplotlib.rcParams[\"lines.linewidth\"] = 2.0\n", "sns.reset_orig()\n", "\n", "# Import tensorboard\n", "%load_ext tensorboard\n", "\n", "# Path to the folder where the datasets are/should be downloaded (e.g. CIFAR10)\n", "DATASET_PATH = os.environ.get(\"PATH_DATASETS\", \"data/\")\n", "# Path to the folder where the pretrained models are saved\n", "CHECKPOINT_PATH = os.environ.get(\"PATH_CHECKPOINT\", \"saved_models/MetaLearning/\")\n", "\n", "# Setting the seed\n", "pl.seed_everything(42)\n", "\n", "# Ensure that all operations are deterministic on GPU (if used) for reproducibility\n", "torch.backends.cudnn.deterministic = True\n", "torch.backends.cudnn.benchmark = False\n", "\n", "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n", "print(\"Device:\", device)"]}, {"cell_type": "markdown", "id": "f3ab8baa", "metadata": {"papermill": {"duration": 0.009264, "end_time": "2025-04-03T19:34:54.264190", "exception": false, "start_time": "2025-04-03T19:34:54.254926", "status": "completed"}, "tags": []}, "source": ["Training the models in this notebook can take between 2 and 8 hours, and the evaluation time of some algorithms is in the span of couples of minutes.\n", "Hence, we download pre-trained models and results below."]}, {"cell_type": "code", "execution_count": 3, "id": "6806656d", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:34:54.284049Z", "iopub.status.busy": "2025-04-03T19:34:54.283654Z", "iopub.status.idle": "2025-04-03T19:34:55.447840Z", "shell.execute_reply": "2025-04-03T19:34:55.446657Z"}, "papermill": {"duration": 1.1761, "end_time": "2025-04-03T19:34:55.449528", "exception": false, "start_time": "2025-04-03T19:34:54.273428", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/ProtoNet.ckpt...\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/ProtoMAML.ckpt...\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/tensorboards/ProtoNet/events.out.tfevents.ProtoNet...\n", "Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/tensorboards/ProtoMAML/events.out.tfevents.ProtoMAML...\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/protomaml_fewshot.json...\n", "Downloading https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/protomaml_svhn_fewshot.json...\n"]}], "source": ["# Github URL where saved models are stored for this tutorial\n", "base_url = \"https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial16/\"\n", "# Files to download\n", "pretrained_files = [\n", " \"ProtoNet.ckpt\",\n", " \"ProtoMAML.ckpt\",\n", " \"tensorboards/ProtoNet/events.out.tfevents.ProtoNet\",\n", " \"tensorboards/ProtoMAML/events.out.tfevents.ProtoMAML\",\n", " \"protomaml_fewshot.json\",\n", " \"protomaml_svhn_fewshot.json\",\n", "]\n", "# Create checkpoint path if it doesn't exist yet\n", "os.makedirs(CHECKPOINT_PATH, exist_ok=True)\n", "\n", "# For each file, check whether it already exists. If not, try downloading it.\n", "for file_name in pretrained_files:\n", " file_path = os.path.join(CHECKPOINT_PATH, file_name)\n", " if \"/\" in file_name:\n", " os.makedirs(file_path.rsplit(\"/\", 1)[0], exist_ok=True)\n", " if not os.path.isfile(file_path):\n", " file_url = base_url + file_name\n", " print(f\"Downloading {file_url}...\")\n", " try:\n", " urllib.request.urlretrieve(file_url, file_path)\n", " except HTTPError as e:\n", " print(\n", " \"Something went wrong. Please try to download the file from the GDrive folder, or contact the author with the full output including the following error:\\n\",\n", " e,\n", " )"]}, {"cell_type": "markdown", "id": "d75a1645", "metadata": {"papermill": {"duration": 0.009471, "end_time": "2025-04-03T19:34:55.468733", "exception": false, "start_time": "2025-04-03T19:34:55.459262", "status": "completed"}, "tags": []}, "source": ["## Few-shot classification\n", "\n", "We start our implementation by discussing the dataset setup.\n", "In this notebook, we will use CIFAR100 which we have already seen in Tutorial 6.\n", "CIFAR100 has 100 classes each with 600 images of size $32\\times 32$ pixels.\n", "Instead of splitting the training, validation and test set over examples, we will split them over classes: we will use 80 classes for training, and 10 for validation and 10 for testing.\n", "Our overall goal is to obtain a model that can distinguish between the 10 test classes with seeing very little examples.\n", "First, let's load the dataset and visualize some examples."]}, {"cell_type": "code", "execution_count": 4, "id": "7233b858", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:34:55.489982Z", "iopub.status.busy": "2025-04-03T19:34:55.488959Z", "iopub.status.idle": "2025-04-03T19:35:02.438492Z", "shell.execute_reply": "2025-04-03T19:35:02.436554Z"}, "papermill": {"duration": 6.962791, "end_time": "2025-04-03T19:35:02.441012", "exception": false, "start_time": "2025-04-03T19:34:55.478221", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Downloading https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz to /__w/11/s/.datasets/cifar-100-python.tar.gz\n"]}, {"name": "stderr", "output_type": "stream", "text": ["\r", " 0%| | 0/169001437 [00:00\n", "\n", "\n"], "text/plain": [""]}, "metadata": {}, "output_type": "display_data"}], "source": ["# Visualize some examples\n", "NUM_IMAGES = 12\n", "cifar_images = [cifar_train_set[np.random.randint(len(cifar_train_set))][0] for idx in range(NUM_IMAGES)]\n", "cifar_images = torch.stack(cifar_images, dim=0)\n", "img_grid = torchvision.utils.make_grid(cifar_images, nrow=6, normalize=True, pad_value=0.9)\n", "img_grid = img_grid.permute(1, 2, 0)\n", "\n", "plt.figure(figsize=(8, 8))\n", "plt.title(\"Image examples of the CIFAR100 dataset\")\n", "plt.imshow(img_grid)\n", "plt.axis(\"off\")\n", "plt.show()\n", "plt.close()"]}, {"cell_type": "markdown", "id": "f950ca64", "metadata": {"papermill": {"duration": 0.027612, "end_time": "2025-04-03T19:35:03.010231", "exception": false, "start_time": "2025-04-03T19:35:02.982619", "status": "completed"}, "tags": []}, "source": ["### Data preprocessing\n", "\n", "Next, we need to prepare the dataset in the training, validation and test split as mentioned before.\n", "The torchvision package gives us the training and test set as two separate dataset objects.\n", "The next code cells will merge the original training and test set, and then create the new train-val-test split."]}, {"cell_type": "code", "execution_count": 6, "id": "16904ddd", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.044121Z", "iopub.status.busy": "2025-04-03T19:35:03.043923Z", "iopub.status.idle": "2025-04-03T19:35:03.083549Z", "shell.execute_reply": "2025-04-03T19:35:03.082343Z"}, "papermill": {"duration": 0.056008, "end_time": "2025-04-03T19:35:03.084829", "exception": false, "start_time": "2025-04-03T19:35:03.028821", "status": "completed"}, "tags": []}, "outputs": [], "source": ["# Merging original training and test set\n", "cifar_all_images = np.concatenate([cifar_train_set.data, cifar_test_set.data], axis=0)\n", "cifar_all_targets = torch.LongTensor(cifar_train_set.targets + cifar_test_set.targets)"]}, {"cell_type": "markdown", "id": "6fd60789", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.018113, "end_time": "2025-04-03T19:35:03.121368", "exception": false, "start_time": "2025-04-03T19:35:03.103255", "status": "completed"}, "tags": []}, "source": ["To have an easier time handling the dataset, we define our own, simple dataset class below.\n", "It takes a set of images, labels/targets, and image transformations, and\n", "returns the corresponding images and labels element-wise."]}, {"cell_type": "code", "execution_count": 7, "id": "a5b97a78", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.154336Z", "iopub.status.busy": "2025-04-03T19:35:03.153936Z", "iopub.status.idle": "2025-04-03T19:35:03.160171Z", "shell.execute_reply": "2025-04-03T19:35:03.159172Z"}, "papermill": {"duration": 0.021873, "end_time": "2025-04-03T19:35:03.161488", "exception": false, "start_time": "2025-04-03T19:35:03.139615", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ImageDataset(data.Dataset):\n", " def __init__(self, imgs, targets, img_transform=None):\n", " \"\"\"\n", " Args:\n", " imgs: Numpy array of shape [N,32,32,3] containing all images.\n", " targets: PyTorch array of shape [N] containing all labels.\n", " img_transform: A torchvision transformation that should be applied\n", " to the images before returning. If none, no transformation\n", " is applied.\n", " \"\"\"\n", " super().__init__()\n", " self.img_transform = img_transform\n", " self.imgs = imgs\n", " self.targets = targets\n", "\n", " def __getitem__(self, idx):\n", " img, target = self.imgs[idx], self.targets[idx]\n", " img = Image.fromarray(img)\n", "\n", " if self.img_transform is not None:\n", " img = self.img_transform(img)\n", "\n", " return img, target\n", "\n", " def __len__(self):\n", " return self.imgs.shape[0]"]}, {"cell_type": "markdown", "id": "7dc8c817", "metadata": {"papermill": {"duration": 0.013117, "end_time": "2025-04-03T19:35:03.187888", "exception": false, "start_time": "2025-04-03T19:35:03.174771", "status": "completed"}, "tags": []}, "source": ["Now, we can create the class splits.\n", "We will assign the classes randomly to training, validation and test, and use a 80%-10%-10% split."]}, {"cell_type": "code", "execution_count": 8, "id": "512186e3", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.215829Z", "iopub.status.busy": "2025-04-03T19:35:03.215473Z", "iopub.status.idle": "2025-04-03T19:35:03.224564Z", "shell.execute_reply": "2025-04-03T19:35:03.223485Z"}, "papermill": {"duration": 0.024859, "end_time": "2025-04-03T19:35:03.225996", "exception": false, "start_time": "2025-04-03T19:35:03.201137", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["Seed set to 0\n"]}], "source": ["pl.seed_everything(0) # Set seed for reproducibility\n", "classes = torch.randperm(100) # Returns random permutation of numbers 0 to 99\n", "train_classes, val_classes, test_classes = classes[:80], classes[80:90], classes[90:]"]}, {"cell_type": "markdown", "id": "abd288e7", "metadata": {"papermill": {"duration": 0.01333, "end_time": "2025-04-03T19:35:03.252739", "exception": false, "start_time": "2025-04-03T19:35:03.239409", "status": "completed"}, "tags": []}, "source": ["To get an intuition of the validation and test classes, we print the class names below:"]}, {"cell_type": "code", "execution_count": 9, "id": "ca6ece69", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.281785Z", "iopub.status.busy": "2025-04-03T19:35:03.280933Z", "iopub.status.idle": "2025-04-03T19:35:03.288678Z", "shell.execute_reply": "2025-04-03T19:35:03.287609Z"}, "papermill": {"duration": 0.024021, "end_time": "2025-04-03T19:35:03.290075", "exception": false, "start_time": "2025-04-03T19:35:03.266054", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Validation classes: ['caterpillar', 'castle', 'skunk', 'ray', 'bus', 'motorcycle', 'keyboard', 'chimpanzee', 'possum', 'tiger']\n", "Test classes: ['kangaroo', 'crocodile', 'butterfly', 'shark', 'forest', 'pickup_truck', 'telephone', 'lion', 'worm', 'mushroom']\n"]}], "source": ["# Printing validation and test classes\n", "idx_to_class = {val: key for key, val in cifar_train_set.class_to_idx.items()}\n", "print(\"Validation classes:\", [idx_to_class[c.item()] for c in val_classes])\n", "print(\"Test classes:\", [idx_to_class[c.item()] for c in test_classes])"]}, {"cell_type": "markdown", "id": "7139e274", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.042815, "end_time": "2025-04-03T19:35:03.346371", "exception": false, "start_time": "2025-04-03T19:35:03.303556", "status": "completed"}, "tags": []}, "source": ["As we can see, the classes have quite some variety and some classes might be easier to distinguish than others.\n", "For instance, in the test classes, 'pickup_truck' is the only vehicle while the classes 'mushroom', 'worm' and 'forest' might be harder to keep apart.\n", "Remember that we want to learn the classification of those ten classes from 80 other classes in our training set, and few examples from the actual test classes.\n", "We will experiment with the number of examples per class.\n", "\n", "Finally, we can create the training, validation and test dataset according to our split above.\n", "For this, we create dataset objects of our previously defined class `ImageDataset`."]}, {"cell_type": "code", "execution_count": 10, "id": "6ca7a69b", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.375219Z", "iopub.status.busy": "2025-04-03T19:35:03.374856Z", "iopub.status.idle": "2025-04-03T19:35:03.381045Z", "shell.execute_reply": "2025-04-03T19:35:03.380012Z"}, "papermill": {"duration": 0.022336, "end_time": "2025-04-03T19:35:03.382433", "exception": false, "start_time": "2025-04-03T19:35:03.360097", "status": "completed"}, "tags": []}, "outputs": [], "source": ["def dataset_from_labels(imgs, targets, class_set, **kwargs):\n", " class_mask = (targets[:, None] == class_set[None, :]).any(dim=-1)\n", " return ImageDataset(imgs=imgs[class_mask], targets=targets[class_mask], **kwargs)"]}, {"cell_type": "markdown", "id": "e727b2c6", "metadata": {"papermill": {"duration": 0.013415, "end_time": "2025-04-03T19:35:03.409336", "exception": false, "start_time": "2025-04-03T19:35:03.395921", "status": "completed"}, "tags": []}, "source": ["As in our experiments before on CIFAR in Tutorial 5, 6 and 9, we normalize the dataset.\n", "Additionally, we use small augmentations during training to prevent overfitting."]}, {"cell_type": "code", "execution_count": 11, "id": "2a58c840", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:03.438219Z", "iopub.status.busy": "2025-04-03T19:35:03.437341Z", "iopub.status.idle": "2025-04-03T19:35:04.567635Z", "shell.execute_reply": "2025-04-03T19:35:04.566112Z"}, "papermill": {"duration": 1.146748, "end_time": "2025-04-03T19:35:04.569567", "exception": false, "start_time": "2025-04-03T19:35:03.422819", "status": "completed"}, "tags": []}, "outputs": [], "source": ["DATA_MEANS = (cifar_train_set.data / 255.0).mean(axis=(0, 1, 2))\n", "DATA_STD = (cifar_train_set.data / 255.0).std(axis=(0, 1, 2))\n", "\n", "test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(DATA_MEANS, DATA_STD)])\n", "# For training, we add some augmentation.\n", "train_transform = transforms.Compose(\n", " [\n", " transforms.RandomHorizontalFlip(),\n", " transforms.RandomResizedCrop((32, 32), scale=(0.8, 1.0), ratio=(0.9, 1.1)),\n", " transforms.ToTensor(),\n", " transforms.Normalize(DATA_MEANS, DATA_STD),\n", " ]\n", ")\n", "\n", "train_set = dataset_from_labels(cifar_all_images, cifar_all_targets, train_classes, img_transform=train_transform)\n", "val_set = dataset_from_labels(cifar_all_images, cifar_all_targets, val_classes, img_transform=test_transform)\n", "test_set = dataset_from_labels(cifar_all_images, cifar_all_targets, test_classes, img_transform=test_transform)"]}, {"cell_type": "markdown", "id": "5cf1f06f", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013398, "end_time": "2025-04-03T19:35:04.596684", "exception": false, "start_time": "2025-04-03T19:35:04.583286", "status": "completed"}, "tags": []}, "source": ["### Data sampling\n", "\n", "The strategy of how to use the available training data for learning few-shot adaptation is crucial in meta-learning.\n", "All three algorithms that we discuss here have a similar idea: simulate few-shot learning during training.\n", "Specifically, at each training step, we randomly select a small number of classes, and sample a small number of examples for each class.\n", "This represents our few-shot training batch, which we also refer to as **support set**.\n", "Additionally, we sample a second set of examples from the same classes, and refer to this batch as **query set**.\n", "Our training objective is to classify the query set correctly from seeing the support set and its corresponding labels.\n", "The main difference between our three methods (ProtoNet, MAML, and Proto-MAML) is in how they use the support set to adapt to the training classes.\n", "\n", "This subsection summarizes the code that is needed to create such training batches.\n", "In PyTorch, we can specify the data sampling procedure by so-called `Sampler` ([documentation](https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler)).\n", "Samplers are iterable objects that return indices in the order in which the data elements should be sampled.\n", "In our previous notebooks, we usually used the option `shuffle=True` in the `data.DataLoader` objects which creates a sampler returning the data indices in a random order.\n", "Here, we focus on samplers that return batches of indices that correspond to support and query set batches.\n", "Below, we implement such a sampler."]}, {"cell_type": "code", "execution_count": 12, "id": "6adc060d", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:04.625385Z", "iopub.status.busy": "2025-04-03T19:35:04.625182Z", "iopub.status.idle": "2025-04-03T19:35:04.636228Z", "shell.execute_reply": "2025-04-03T19:35:04.635248Z"}, "papermill": {"duration": 0.02732, "end_time": "2025-04-03T19:35:04.637546", "exception": false, "start_time": "2025-04-03T19:35:04.610226", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class FewShotBatchSampler:\n", " def __init__(self, dataset_targets, N_way, K_shot, include_query=False, shuffle=True, shuffle_once=False):\n", " \"\"\"FewShot Batch Sampler.\n", "\n", " Args:\n", " dataset_targets: PyTorch tensor of the labels of the data elements.\n", " N_way: Number of classes to sample per batch.\n", " K_shot: Number of examples to sample per class in the batch.\n", " include_query: If True, returns batch of size N_way*K_shot*2, which\n", " can be split into support and query set. Simplifies\n", " the implementation of sampling the same classes but\n", " distinct examples for support and query set.\n", " shuffle: If True, examples and classes are newly shuffled in each\n", " iteration (for training)\n", " shuffle_once: If True, examples and classes are shuffled once in\n", " the beginning, but kept constant across iterations\n", " (for validation)\n", "\n", " \"\"\"\n", " super().__init__()\n", " self.dataset_targets = dataset_targets\n", " self.N_way = N_way\n", " self.K_shot = K_shot\n", " self.shuffle = shuffle\n", " self.include_query = include_query\n", " if self.include_query:\n", " self.K_shot *= 2\n", " self.batch_size = self.N_way * self.K_shot # Number of overall images per batch\n", "\n", " # Organize examples by class\n", " self.classes = torch.unique(self.dataset_targets).tolist()\n", " self.num_classes = len(self.classes)\n", " self.indices_per_class = {}\n", " self.batches_per_class = {} # Number of K-shot batches that each class can provide\n", " for c in self.classes:\n", " self.indices_per_class[c] = torch.where(self.dataset_targets == c)[0]\n", " self.batches_per_class[c] = self.indices_per_class[c].shape[0] // self.K_shot\n", "\n", " # Create a list of classes from which we select the N classes per batch\n", " self.iterations = sum(self.batches_per_class.values()) // self.N_way\n", " self.class_list = [c for c in self.classes for _ in range(self.batches_per_class[c])]\n", " if shuffle_once or self.shuffle:\n", " self.shuffle_data()\n", " else:\n", " # For testing, we iterate over classes instead of shuffling them\n", " sort_idxs = [\n", " i + p * self.num_classes for i, c in enumerate(self.classes) for p in range(self.batches_per_class[c])\n", " ]\n", " self.class_list = np.array(self.class_list)[np.argsort(sort_idxs)].tolist()\n", "\n", " def shuffle_data(self):\n", " # Shuffle the examples per class\n", " for c in self.classes:\n", " perm = torch.randperm(self.indices_per_class[c].shape[0])\n", " self.indices_per_class[c] = self.indices_per_class[c][perm]\n", " # Shuffle the class list from which we sample. Note that this way of shuffling\n", " # does not prevent to choose the same class twice in a batch. However, for\n", " # training and validation, this is not a problem.\n", " random.shuffle(self.class_list)\n", "\n", " def __iter__(self):\n", " # Shuffle data\n", " if self.shuffle:\n", " self.shuffle_data()\n", "\n", " # Sample few-shot batches\n", " start_index = defaultdict(int)\n", " for it in range(self.iterations):\n", " class_batch = self.class_list[it * self.N_way : (it + 1) * self.N_way] # Select N classes for the batch\n", " index_batch = []\n", " for c in class_batch: # For each class, select the next K examples and add them to the batch\n", " index_batch.extend(self.indices_per_class[c][start_index[c] : start_index[c] + self.K_shot])\n", " start_index[c] += self.K_shot\n", " if self.include_query: # If we return support+query set, sort them so that they are easy to split\n", " index_batch = index_batch[::2] + index_batch[1::2]\n", " yield index_batch\n", "\n", " def __len__(self):\n", " return self.iterations"]}, {"cell_type": "markdown", "id": "99bb3f79", "metadata": {"papermill": {"duration": 0.013412, "end_time": "2025-04-03T19:35:04.664501", "exception": false, "start_time": "2025-04-03T19:35:04.651089", "status": "completed"}, "tags": []}, "source": ["Now, we can create our intended data loaders by passing an object of `FewShotBatchSampler` as `batch_sampler=...` input to the PyTorch data loader object.\n", "For our experiments, we will use a 5-class 4-shot training setting.\n", "This means that each support set contains 5 classes with 4 examples each, i.e., 20 images overall.\n", "Usually, it is good to keep the number of shots equal to the number that you aim to test on.\n", "However, we will experiment later with different number of shots, and hence, we pick 4 as a compromise for now.\n", "To get the best performing model, it is recommended to consider the\n", "number of training shots as hyperparameter in a grid search."]}, {"cell_type": "code", "execution_count": 13, "id": "df85b558", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:04.692478Z", "iopub.status.busy": "2025-04-03T19:35:04.692252Z", "iopub.status.idle": "2025-04-03T19:35:04.726396Z", "shell.execute_reply": "2025-04-03T19:35:04.725401Z"}, "papermill": {"duration": 0.049774, "end_time": "2025-04-03T19:35:04.727749", "exception": false, "start_time": "2025-04-03T19:35:04.677975", "status": "completed"}, "tags": []}, "outputs": [], "source": ["N_WAY = 5\n", "K_SHOT = 4\n", "train_data_loader = data.DataLoader(\n", " train_set,\n", " batch_sampler=FewShotBatchSampler(train_set.targets, include_query=True, N_way=N_WAY, K_shot=K_SHOT, shuffle=True),\n", " num_workers=4,\n", ")\n", "val_data_loader = data.DataLoader(\n", " val_set,\n", " batch_sampler=FewShotBatchSampler(\n", " val_set.targets, include_query=True, N_way=N_WAY, K_shot=K_SHOT, shuffle=False, shuffle_once=True\n", " ),\n", " num_workers=4,\n", ")"]}, {"cell_type": "markdown", "id": "3935d548", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013461, "end_time": "2025-04-03T19:35:04.754841", "exception": false, "start_time": "2025-04-03T19:35:04.741380", "status": "completed"}, "tags": []}, "source": ["For simplicity, we implemented the sampling of a support and query set as sampling a support set with twice the number of examples.\n", "After sampling a batch from the data loader, we need to split it into a support and query set.\n", "We can summarize this step in the following function:"]}, {"cell_type": "code", "execution_count": 14, "id": "c1f8b80b", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:04.782961Z", "iopub.status.busy": "2025-04-03T19:35:04.782732Z", "iopub.status.idle": "2025-04-03T19:35:04.787385Z", "shell.execute_reply": "2025-04-03T19:35:04.786395Z"}, "papermill": {"duration": 0.020259, "end_time": "2025-04-03T19:35:04.788659", "exception": false, "start_time": "2025-04-03T19:35:04.768400", "status": "completed"}, "tags": []}, "outputs": [], "source": ["def split_batch(imgs, targets):\n", " support_imgs, query_imgs = imgs.chunk(2, dim=0)\n", " support_targets, query_targets = targets.chunk(2, dim=0)\n", " return support_imgs, query_imgs, support_targets, query_targets"]}, {"cell_type": "markdown", "id": "64d20081", "metadata": {"papermill": {"duration": 0.013473, "end_time": "2025-04-03T19:35:04.815722", "exception": false, "start_time": "2025-04-03T19:35:04.802249", "status": "completed"}, "tags": []}, "source": ["Finally, to ensure that our implementation of the data sampling process is correct, we can sample a batch and visualize its support and query set.\n", "What we would like to see is that the support and query set have the same classes, but distinct examples."]}, {"cell_type": "code", "execution_count": 15, "id": "0ffbfb0e", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:04.843982Z", "iopub.status.busy": "2025-04-03T19:35:04.843616Z", "iopub.status.idle": "2025-04-03T19:35:05.180823Z", "shell.execute_reply": "2025-04-03T19:35:05.179756Z"}, "papermill": {"duration": 0.353014, "end_time": "2025-04-03T19:35:05.182266", "exception": false, "start_time": "2025-04-03T19:35:04.829252", "status": "completed"}, "tags": []}, "outputs": [], "source": ["imgs, targets = next(iter(val_data_loader)) # We use the validation set since it does not apply augmentations\n", "support_imgs, query_imgs, _, _ = split_batch(imgs, targets)\n", "support_grid = torchvision.utils.make_grid(support_imgs, nrow=K_SHOT, normalize=True, pad_value=0.9)\n", "support_grid = support_grid.permute(1, 2, 0)\n", "query_grid = torchvision.utils.make_grid(query_imgs, nrow=K_SHOT, normalize=True, pad_value=0.9)\n", "query_grid = query_grid.permute(1, 2, 0)\n", "\n", "fig, ax = plt.subplots(1, 2, figsize=(8, 5))\n", "ax[0].imshow(support_grid)\n", "ax[0].set_title(\"Support set\")\n", "ax[0].axis(\"off\")\n", "ax[1].imshow(query_grid)\n", "ax[1].set_title(\"Query set\")\n", "ax[1].axis(\"off\")\n", "fig.suptitle(\"Few Shot Batch\", weight=\"bold\")\n", "fig.show()\n", "plt.close(fig)"]}, {"cell_type": "markdown", "id": "865ba802", "metadata": {"papermill": {"duration": 0.013939, "end_time": "2025-04-03T19:35:05.210379", "exception": false, "start_time": "2025-04-03T19:35:05.196440", "status": "completed"}, "tags": []}, "source": ["As we can see, the support and query set have the same five classes, but different examples.\n", "The models will be tasked to classify the examples in the query set by learning from the support set and its labels.\n", "With the data sampling in place, we can now start to implement our first meta-learning model: Prototypical Networks."]}, {"cell_type": "markdown", "id": "6874c15b", "metadata": {"papermill": {"duration": 0.013528, "end_time": "2025-04-03T19:35:05.237585", "exception": false, "start_time": "2025-04-03T19:35:05.224057", "status": "completed"}, "tags": []}, "source": ["## Prototypical Networks\n", "\n", "
"]}, {"cell_type": "markdown", "id": "44f9d0e6", "metadata": {"papermill": {"duration": 0.013597, "end_time": "2025-04-03T19:35:05.264707", "exception": false, "start_time": "2025-04-03T19:35:05.251110", "status": "completed"}, "tags": []}, "source": ["The Prototypical Network, or ProtoNet for short, is a metric-based meta-learning algorithm which operates similar to a nearest neighbor classification.\n", "Metric-based meta-learning methods classify a new example $\\mathbf{x}$ based on some distance function $d_{\\varphi}$ between $x$ and all elements in the support set.\n", "ProtoNets implements this idea with the concept of prototypes in a learned feature space.\n", "First, ProtoNet uses an embedding function $f_{\\theta}$ to encode each input in the support set into a $L$-dimensional feature vector.\n", "Next, for each class $c$, we collect the feature vectors of all examples with label $c$, and average their feature vectors.\n", "Formally, we can define this as:\n", "\n", "$$\\mathbf{v}_c=\\frac{1}{|S_c|}\\sum_{(\\mathbf{x}_i,y_i)\\in S_c}f_{\\theta}(\\mathbf{x}_i)$$\n", "\n", "where $S_c$ is the part of the support set $S$ for which $y_i=c$, and $\\mathbf{v}_c$ represents the _prototype_ of class $c$.\n", "The prototype calculation is visualized below for a 2-dimensional feature space and 3 classes (Figure credit - [Snell et al.](https://arxiv.org/abs/1703.05175)).\n", "The colored dots represent encoded support elements with color-corresponding class label, and the black dots next to the class label are the averaged prototypes.\n", "\n", "
\n", "\n", "Based on these prototypes, we want to classify a new example.\n", "Remember that since we want to learn the encoding function $f_{\\theta}$, this classification must be differentiable and hence, we need to define a probability distribution across classes.\n", "For this, we will make use of the distance function $d_{\\varphi}$: the closer a new example $\\mathbf{x}$ is to a prototype $\\mathbf{v}_c$, the higher the probability for $\\mathbf{x}$ belonging to class $c$.\n", "Formally, we can simply use a softmax over the distances of $\\mathbf{x}$ to all class prototypes:\n", "\n", "$$p(y=c\\vert\\mathbf{x})=\\text{softmax}(-d_{\\varphi}(f_{\\theta}(\\mathbf{x}), \\mathbf{v}_c))=\\frac{\\exp\\left(-d_{\\varphi}(f_{\\theta}(\\mathbf{x}), \\mathbf{v}_c)\\right)}{\\sum_{c'\\in \\mathcal{C}}\\exp\\left(-d_{\\varphi}(f_{\\theta}(\\mathbf{x}), \\mathbf{v}_{c'})\\right)}$$\n", "\n", "Note that the negative sign is necessary since we want to increase the probability for close-by vectors and have a low probability for distant vectors.\n", "We train the network $f_{\\theta}$ based on the cross entropy error of the training query set examples.\n", "Thereby, the gradient flows through both the prototypes $\\mathbf{v}_c$ and the query set encodings $f_{\\theta}(\\mathbf{x})$.\n", "For the distance function $d_{\\varphi}$, we can choose any function as long as it is differentiable with respect to both of its inputs.\n", "The most common function, which we also use here, is the squared\n", "euclidean distance, but there has been several works on different\n", "distance functions as well."]}, {"cell_type": "markdown", "id": "e7de47ca", "metadata": {"papermill": {"duration": 0.014283, "end_time": "2025-04-03T19:35:05.292538", "exception": false, "start_time": "2025-04-03T19:35:05.278255", "status": "completed"}, "tags": []}, "source": ["### ProtoNet implementation"]}, {"cell_type": "markdown", "id": "b3cde232", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013536, "end_time": "2025-04-03T19:35:05.319858", "exception": false, "start_time": "2025-04-03T19:35:05.306322", "status": "completed"}, "tags": []}, "source": ["Now that we know how a ProtoNet works in principle, let's look at how we can apply to our specific problem of few-shot image classification, and implement it below.\n", "First, we need to define the encoder function $f_{\\theta}$.\n", "Since we work with CIFAR images, we can take a look back at Tutorial 5 where we compared common Computer Vision architectures, and choose one of the best performing ones.\n", "Here, we go with a DenseNet since it is in general more parameter efficient than ResNet.\n", "Luckily, we do not need to implement DenseNet ourselves again and can rely on torchvision's model package instead.\n", "We use common hyperparameters of 64 initial feature channels, add 32 per block, and use a bottleneck size of 64 (i.e. 2 times the growth rate).\n", "We use 4 stages of 6 layers each, which results in overall about 1 million parameters.\n", "Note that the torchvision package assumes that the last layer is used for classification and hence calls its output size `num_classes`.\n", "However, we can instead just use it as the feature space of ProtoNet, and choose an arbitrary dimensionality.\n", "We will use the same network for other algorithms in this notebook to ensure a fair comparison."]}, {"cell_type": "code", "execution_count": 16, "id": "723b50aa", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:05.348860Z", "iopub.status.busy": "2025-04-03T19:35:05.348463Z", "iopub.status.idle": "2025-04-03T19:35:05.355180Z", "shell.execute_reply": "2025-04-03T19:35:05.354004Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.023169, "end_time": "2025-04-03T19:35:05.356596", "exception": false, "start_time": "2025-04-03T19:35:05.333427", "status": "completed"}, "tags": []}, "outputs": [], "source": ["def get_convnet(output_size):\n", " convnet = torchvision.models.DenseNet(\n", " growth_rate=32,\n", " block_config=(6, 6, 6, 6),\n", " bn_size=2,\n", " num_init_features=64,\n", " num_classes=output_size, # Output dimensionality\n", " )\n", " return convnet"]}, {"cell_type": "markdown", "id": "0e91710f", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013617, "end_time": "2025-04-03T19:35:05.383862", "exception": false, "start_time": "2025-04-03T19:35:05.370245", "status": "completed"}, "tags": []}, "source": ["Next, we can look at implementing ProtoNet.\n", "We will define it as PyTorch Lightning module to use all functionalities of PyTorch Lightning.\n", "The first step during training is to encode all images in a batch with our network.\n", "Next, we calculate the class prototypes from the support set (function `calculate_prototypes`), and classify the query set examples according to the prototypes (function `classify_feats`).\n", "Keep in mind that we use the data sampling described before, such that the support and query set are stacked together in the batch.\n", "Thus, we use our previously defined function `split_batch` to split them apart.\n", "The full code can be found below."]}, {"cell_type": "code", "execution_count": 17, "id": "73ad6acc", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:05.412851Z", "iopub.status.busy": "2025-04-03T19:35:05.412484Z", "iopub.status.idle": "2025-04-03T19:35:05.429220Z", "shell.execute_reply": "2025-04-03T19:35:05.428106Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.033054, "end_time": "2025-04-03T19:35:05.430636", "exception": false, "start_time": "2025-04-03T19:35:05.397582", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ProtoNet(pl.LightningModule):\n", " def __init__(self, proto_dim, lr):\n", " \"\"\"ProtoNet.\n", "\n", " Args:\n", " proto_dim: Dimensionality of prototype feature space\n", " lr: Learning rate of Adam optimizer\n", "\n", " \"\"\"\n", " super().__init__()\n", " self.save_hyperparameters()\n", " self.model = get_convnet(output_size=self.hparams.proto_dim)\n", "\n", " def configure_optimizers(self):\n", " optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr)\n", " scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[140, 180], gamma=0.1)\n", " return [optimizer], [scheduler]\n", "\n", " @staticmethod\n", " def calculate_prototypes(features, targets):\n", " # Given a stack of features vectors and labels, return class prototypes\n", " # features - shape [N, proto_dim], targets - shape [N]\n", " classes, _ = torch.unique(targets).sort() # Determine which classes we have\n", " prototypes = []\n", " for c in classes:\n", " p = features[torch.where(targets == c)[0]].mean(dim=0) # Average class feature vectors\n", " prototypes.append(p)\n", " prototypes = torch.stack(prototypes, dim=0)\n", " # Return the 'classes' tensor to know which prototype belongs to which class\n", " return prototypes, classes\n", "\n", " def classify_feats(self, prototypes, classes, feats, targets):\n", " # Classify new examples with prototypes and return classification error\n", " dist = torch.pow(prototypes[None, :] - feats[:, None], 2).sum(dim=2) # Squared euclidean distance\n", " preds = F.log_softmax(-dist, dim=1)\n", " labels = (classes[None, :] == targets[:, None]).long().argmax(dim=-1)\n", " acc = (preds.argmax(dim=1) == labels).float().mean()\n", " return preds, labels, acc\n", "\n", " def calculate_loss(self, batch, mode):\n", " # Determine training loss for a given support and query set\n", " imgs, targets = batch\n", " features = self.model(imgs) # Encode all images of support and query set\n", " support_feats, query_feats, support_targets, query_targets = split_batch(features, targets)\n", " prototypes, classes = ProtoNet.calculate_prototypes(support_feats, support_targets)\n", " preds, labels, acc = self.classify_feats(prototypes, classes, query_feats, query_targets)\n", " loss = F.cross_entropy(preds, labels)\n", "\n", " self.log(f\"{mode}_loss\", loss)\n", " self.log(f\"{mode}_acc\", acc)\n", " return loss\n", "\n", " def training_step(self, batch, batch_idx):\n", " return self.calculate_loss(batch, mode=\"train\")\n", "\n", " def validation_step(self, batch, batch_idx):\n", " self.calculate_loss(batch, mode=\"val\")"]}, {"cell_type": "markdown", "id": "1f70fab9", "metadata": {"papermill": {"duration": 0.013715, "end_time": "2025-04-03T19:35:05.458181", "exception": false, "start_time": "2025-04-03T19:35:05.444466", "status": "completed"}, "tags": []}, "source": ["For validation, we use the same principle as training and sample support and query sets from the hold-out 10 classes.\n", "However, this gives us noisy scores depending on which query sets are chosen to which support sets.\n", "This is why we will use a different strategy during testing.\n", "For validation, our training strategy is sufficient since it is much\n", "faster than testing, and gives a good estimate of the training\n", "generalization as long as we keep the support-query sets constant across\n", "validation iterations."]}, {"cell_type": "markdown", "id": "b96ca727", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013641, "end_time": "2025-04-03T19:35:05.485461", "exception": false, "start_time": "2025-04-03T19:35:05.471820", "status": "completed"}, "tags": []}, "source": ["### Training\n", "\n", "After implementing the model, we can already start training it.\n", "We use our common PyTorch Lightning training function, and train the model for 200 epochs.\n", "The training function takes `model_class` as input argument, i.e. the\n", "PyTorch Lightning module class that should be trained, since we will\n", "reuse this function for other algorithms as well."]}, {"cell_type": "code", "execution_count": 18, "id": "e00938e3", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:05.515065Z", "iopub.status.busy": "2025-04-03T19:35:05.513945Z", "iopub.status.idle": "2025-04-03T19:35:05.523361Z", "shell.execute_reply": "2025-04-03T19:35:05.522232Z"}, "papermill": {"duration": 0.025663, "end_time": "2025-04-03T19:35:05.524823", "exception": false, "start_time": "2025-04-03T19:35:05.499160", "status": "completed"}, "tags": []}, "outputs": [], "source": ["def train_model(model_class, train_loader, val_loader, **kwargs):\n", " trainer = pl.Trainer(\n", " default_root_dir=os.path.join(CHECKPOINT_PATH, model_class.__name__),\n", " accelerator=\"auto\",\n", " devices=1,\n", " max_epochs=200,\n", " callbacks=[\n", " ModelCheckpoint(save_weights_only=True, mode=\"max\", monitor=\"val_acc\"),\n", " LearningRateMonitor(\"epoch\"),\n", " ],\n", " enable_progress_bar=False,\n", " )\n", " trainer.logger._default_hp_metric = None\n", "\n", " # Check whether pretrained model exists. If yes, load it and skip training\n", " pretrained_filename = os.path.join(CHECKPOINT_PATH, model_class.__name__ + \".ckpt\")\n", " if os.path.isfile(pretrained_filename):\n", " print(f\"Found pretrained model at {pretrained_filename}, loading...\")\n", " # Automatically loads the model with the saved hyperparameters\n", " model = model_class.load_from_checkpoint(pretrained_filename)\n", " else:\n", " pl.seed_everything(42) # To be reproducible\n", " model = model_class(**kwargs)\n", " trainer.fit(model, train_loader, val_loader)\n", " model = model_class.load_from_checkpoint(\n", " trainer.checkpoint_callback.best_model_path\n", " ) # Load best checkpoint after training\n", "\n", " return model"]}, {"cell_type": "markdown", "id": "447f6b30", "metadata": {"papermill": {"duration": 0.01364, "end_time": "2025-04-03T19:35:05.552293", "exception": false, "start_time": "2025-04-03T19:35:05.538653", "status": "completed"}, "tags": []}, "source": ["Below is the training call for our ProtoNet.\n", "We use a 64-dimensional feature space.\n", "Larger feature spaces showed to give noisier results since the squared euclidean distance becomes proportionally larger in expectation, and smaller feature spaces might not allow for enough flexibility.\n", "We recommend to load the pre-trained model here at first, but feel free\n", "to play around with the hyperparameters yourself."]}, {"cell_type": "code", "execution_count": 19, "id": "d08b68a6", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:05.581204Z", "iopub.status.busy": "2025-04-03T19:35:05.580841Z", "iopub.status.idle": "2025-04-03T19:35:06.335451Z", "shell.execute_reply": "2025-04-03T19:35:06.334286Z"}, "papermill": {"duration": 0.771558, "end_time": "2025-04-03T19:35:06.337621", "exception": false, "start_time": "2025-04-03T19:35:05.566063", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["GPU available: True (cuda), used: True\n"]}, {"name": "stderr", "output_type": "stream", "text": ["TPU available: False, using: 0 TPU cores\n"]}, {"name": "stderr", "output_type": "stream", "text": ["HPU available: False, using: 0 HPUs\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Found pretrained model at saved_models/MetaLearning/ProtoNet.ckpt, loading...\n"]}, {"name": "stderr", "output_type": "stream", "text": ["Lightning automatically upgraded your loaded checkpoint from v1.3.4 to v2.4.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint saved_models/MetaLearning/ProtoNet.ckpt`\n"]}], "source": ["protonet_model = train_model(\n", " ProtoNet, proto_dim=64, lr=2e-4, train_loader=train_data_loader, val_loader=val_data_loader\n", ")"]}, {"cell_type": "markdown", "id": "545ece26", "metadata": {"papermill": {"duration": 0.025979, "end_time": "2025-04-03T19:35:06.394973", "exception": false, "start_time": "2025-04-03T19:35:06.368994", "status": "completed"}, "tags": []}, "source": ["We can also take a closer look at the TensorBoard below."]}, {"cell_type": "code", "execution_count": 20, "id": "8aa16039", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:06.430691Z", "iopub.status.busy": "2025-04-03T19:35:06.430494Z", "iopub.status.idle": "2025-04-03T19:35:06.434384Z", "shell.execute_reply": "2025-04-03T19:35:06.433389Z"}, "papermill": {"duration": 0.021429, "end_time": "2025-04-03T19:35:06.435596", "exception": false, "start_time": "2025-04-03T19:35:06.414167", "status": "completed"}, "tags": []}, "outputs": [], "source": ["# Opens tensorboard in notebook. Adjust the path to your CHECKPOINT_PATH if needed\n", "# %tensorboard --logdir ../saved_models/tutorial16/tensorboards/ProtoNet/"]}, {"cell_type": "markdown", "id": "df9627d0", "metadata": {"papermill": {"duration": 0.013956, "end_time": "2025-04-03T19:35:06.463640", "exception": false, "start_time": "2025-04-03T19:35:06.449684", "status": "completed"}, "tags": []}, "source": ["
\n", "\n", "In contrast to standard supervised learning, we see that ProtoNet does not overfit as much as we would expect.\n", "The validation accuracy is of course lower than the average training, but the training loss does not stick close to zero.\n", "This is because no training batch is as the other, and we also mix new examples in the support set and query set.\n", "This gives us slightly different prototypes in every iteration, and makes it harder for the network to fully overfit."]}, {"cell_type": "markdown", "id": "0ae78385", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.013973, "end_time": "2025-04-03T19:35:06.491609", "exception": false, "start_time": "2025-04-03T19:35:06.477636", "status": "completed"}, "tags": []}, "source": ["### Testing\n", "\n", "Our goal of meta-learning is to obtain a model that can quickly adapt to a new task, or in this case, new classes to distinguish between.\n", "To test this, we will use our trained ProtoNet and adapt it to the 10 test classes.\n", "Thereby, we pick $k$ examples per class from which we determine the prototypes, and test the classification accuracy on all other examples.\n", "This can be seen as using the $k$ examples per class as support set, and the rest of the dataset as a query set.\n", "We iterate through the dataset such that each example has been once included in a support set.\n", "The average performance over all support sets tells us how well we can expect ProtoNet to perform when seeing only $k$ examples per class.\n", "During training, we used $k=4$.\n", "In testing, we will experiment with $k=\\{2,4,8,16,32\\}$ to get a better sense of how $k$ influences the results.\n", "We would expect that we achieve higher accuracies the more examples we have in the support set, but we don't know how it scales.\n", "Hence, let's first implement a function that executes the testing procedure for a given $k$:"]}, {"cell_type": "code", "execution_count": 21, "id": "25954ba7", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:06.521313Z", "iopub.status.busy": "2025-04-03T19:35:06.520578Z", "iopub.status.idle": "2025-04-03T19:35:06.537447Z", "shell.execute_reply": "2025-04-03T19:35:06.536289Z"}, "papermill": {"duration": 0.033083, "end_time": "2025-04-03T19:35:06.538679", "exception": false, "start_time": "2025-04-03T19:35:06.505596", "status": "completed"}, "tags": []}, "outputs": [], "source": ["@torch.no_grad()\n", "def test_proto_net(model, dataset, data_feats=None, k_shot=4):\n", " \"\"\"Test proto net.\n", "\n", " Args:\n", " model: Pretrained ProtoNet model\n", " dataset: The dataset on which the test should be performed.\n", " Should be instance of ImageDataset\n", " data_feats: The encoded features of all images in the dataset.\n", " If None, they will be newly calculated, and returned\n", " for later usage.\n", " k_shot: Number of examples per class in the support set.\n", "\n", " \"\"\"\n", " model = model.to(device)\n", " model.eval()\n", " num_classes = dataset.targets.unique().shape[0]\n", " exmps_per_class = dataset.targets.shape[0] // num_classes # We assume uniform example distribution here\n", "\n", " # The encoder network remains unchanged across k-shot settings. Hence, we only need\n", " # to extract the features for all images once.\n", " if data_feats is None:\n", " # Dataset preparation\n", " dataloader = data.DataLoader(dataset, batch_size=128, num_workers=4, shuffle=False, drop_last=False)\n", "\n", " img_features = []\n", " img_targets = []\n", " for imgs, targets in tqdm(dataloader, \"Extracting image features\", leave=False):\n", " imgs = imgs.to(device)\n", " feats = model.model(imgs)\n", " img_features.append(feats.detach().cpu())\n", " img_targets.append(targets)\n", " img_features = torch.cat(img_features, dim=0)\n", " img_targets = torch.cat(img_targets, dim=0)\n", " # Sort by classes, so that we obtain tensors of shape [num_classes, exmps_per_class, ...]\n", " # Makes it easier to process later\n", " img_targets, sort_idx = img_targets.sort()\n", " img_targets = img_targets.reshape(num_classes, exmps_per_class).transpose(0, 1)\n", " img_features = img_features[sort_idx].reshape(num_classes, exmps_per_class, -1).transpose(0, 1)\n", " else:\n", " img_features, img_targets = data_feats\n", "\n", " # We iterate through the full dataset in two manners. First, to select the k-shot batch.\n", " # Second, the evaluate the model on all other examples\n", " accuracies = []\n", " for k_idx in tqdm(range(0, img_features.shape[0], k_shot), \"Evaluating prototype classification\", leave=False):\n", " # Select support set and calculate prototypes\n", " k_img_feats = img_features[k_idx : k_idx + k_shot].flatten(0, 1)\n", " k_targets = img_targets[k_idx : k_idx + k_shot].flatten(0, 1)\n", " prototypes, proto_classes = model.calculate_prototypes(k_img_feats, k_targets)\n", " # Evaluate accuracy on the rest of the dataset\n", " batch_acc = 0\n", " for e_idx in range(0, img_features.shape[0], k_shot):\n", " if k_idx == e_idx: # Do not evaluate on the support set examples\n", " continue\n", " e_img_feats = img_features[e_idx : e_idx + k_shot].flatten(0, 1)\n", " e_targets = img_targets[e_idx : e_idx + k_shot].flatten(0, 1)\n", " _, _, acc = model.classify_feats(prototypes, proto_classes, e_img_feats, e_targets)\n", " batch_acc += acc.item()\n", " batch_acc /= img_features.shape[0] // k_shot - 1\n", " accuracies.append(batch_acc)\n", "\n", " return (mean(accuracies), stdev(accuracies)), (img_features, img_targets)"]}, {"cell_type": "markdown", "id": "dcd83de3", "metadata": {"papermill": {"duration": 0.014101, "end_time": "2025-04-03T19:35:06.566900", "exception": false, "start_time": "2025-04-03T19:35:06.552799", "status": "completed"}, "tags": []}, "source": ["Testing ProtoNet is relatively quick if we have processed all images once. Hence, we can do in this notebook:"]}, {"cell_type": "code", "execution_count": 22, "id": "41dbf139", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:06.596678Z", "iopub.status.busy": "2025-04-03T19:35:06.596331Z", "iopub.status.idle": "2025-04-03T19:35:25.411809Z", "shell.execute_reply": "2025-04-03T19:35:25.410699Z"}, "papermill": {"duration": 18.831737, "end_time": "2025-04-03T19:35:25.412953", "exception": false, "start_time": "2025-04-03T19:35:06.581216", "status": "completed"}, "tags": []}, "outputs": [{"data": {"application/vnd.jupyter.widget-view+json": {"model_id": "1516abb90d254a4c94b9885ce76994fd", "version_major": 2, "version_minor": 0}, "text/plain": ["Extracting image features: 0%| | 0/47 [00:00\n", "\n", "\n"], "text/plain": [""]}, "metadata": {}, "output_type": "display_data"}], "source": ["ax = plot_few_shot(protonet_accuracies, name=\"ProtoNet\", color=\"C1\")\n", "plt.show()\n", "plt.close()"]}, {"cell_type": "markdown", "id": "e53b6159", "metadata": {"papermill": {"duration": 0.026965, "end_time": "2025-04-03T19:35:45.885235", "exception": false, "start_time": "2025-04-03T19:35:45.858270", "status": "completed"}, "tags": []}, "source": ["As we initially expected, the performance of ProtoNet indeed increases the more samples we have.\n", "However, even with just two samples per class, we classify almost half of the images correctly, which is well above random accuracy (10%).\n", "The curve shows an exponentially dampend trend, meaning that adding 2 extra examples to $k=2$ has a much higher impact than adding 2 extra samples if we already have $k=16$.\n", "Nonetheless, we can say that ProtoNet adapts fairly well to new classes."]}, {"cell_type": "markdown", "id": "8f5114b1", "metadata": {"papermill": {"duration": 0.016106, "end_time": "2025-04-03T19:35:45.922413", "exception": false, "start_time": "2025-04-03T19:35:45.906307", "status": "completed"}, "tags": []}, "source": ["## MAML and ProtoMAML\n", "\n", "
"]}, {"cell_type": "markdown", "id": "b99d8fb7", "metadata": {"papermill": {"duration": 0.016242, "end_time": "2025-04-03T19:35:45.954960", "exception": false, "start_time": "2025-04-03T19:35:45.938718", "status": "completed"}, "tags": []}, "source": ["The second meta-learning algorithm we will look at is MAML, short for Model-Agnostic Meta-Learning.\n", "MAML is an optimization-based meta-learning algorithm, which means that it tries to adjust the standard optimization procedure to a few-shot setting.\n", "The idea of MAML is relatively simple: given a model, support and query set during training, we optimize the model for $m$ steps on the support set, and evaluate the gradients of the query loss with respect to the original model's parameters.\n", "For the same model, we do it for a few different support-query sets and accumulate the gradients.\n", "This results in learning a model that provides a good initialization for being quickly adapted to the training tasks.\n", "If we denote the model parameters with $\\theta$, we can visualize the procedure as follows (Figure credit - [Finn et al. ](http://proceedings.mlr.press/v70/finn17a.html)).\n", "\n", "
"]}, {"cell_type": "markdown", "id": "300779ff", "metadata": {"papermill": {"duration": 0.016296, "end_time": "2025-04-03T19:35:45.987606", "exception": false, "start_time": "2025-04-03T19:35:45.971310", "status": "completed"}, "tags": []}, "source": ["The full algorithm of MAML is therefore as follows.\n", "At each training step, we sample a batch of tasks, i.e., a batch of support-query set pairs.\n", "For each task $\\mathcal{T}_i$, we optimize a model $f_{\\theta}$ on the support set via SGD, and denote this model as $f_{\\theta_i'}$.\n", "We refer to this optimization as _inner loop_.\n", "Using this new model, we calculate the gradients of the original parameters, $\\theta$, with respect to the query loss on $f_{\\theta_i'}$.\n", "These gradients are accumulated over all tasks, and used to update $\\theta$.\n", "This is called _outer loop_ since we iterate over tasks.\n", "The full MAML algorithm is summarized below (Figure credit - [Finn et al. ](http://proceedings.mlr.press/v70/finn17a.html)).\n", "\n", "
"]}, {"cell_type": "markdown", "id": "c6912ea4", "metadata": {"papermill": {"duration": 0.016272, "end_time": "2025-04-03T19:35:46.020087", "exception": false, "start_time": "2025-04-03T19:35:46.003815", "status": "completed"}, "tags": []}, "source": ["To obtain gradients for the initial parameters $\\theta$ from the optimized model $f_{\\theta_i'}$, we actually need second-order gradients, i.e. gradients of gradients, as the support set gradients depend on $\\theta$ as well.\n", "This makes MAML computationally expensive, especially when using multiple inner loop steps.\n", "A simpler, yet almost equally well performing alternative is First-Order MAML (FOMAML) which only uses first-order gradients.\n", "This means that the second-order gradients are ignored, and we can calculate the outer loop gradients (line 10 in algorithm 2) simply by calculating the gradients with respect to $\\theta_i'$, and use those as update to $\\theta$.\n", "Hence, the new update rule becomes:\n", "$$\\theta\\leftarrow\\theta-\\beta\\sum_{\\mathcal{T}_i\\sim p(\\mathcal{T})}\\nabla_{\\theta_i'}\\mathcal{L}_{\\mathcal{T}_i}(f_{\\theta_i'})$$\n", "Note the change of $\\theta$ to $\\theta_i'$ for $\\nabla$."]}, {"cell_type": "markdown", "id": "825fa4bc", "metadata": {"papermill": {"duration": 0.016206, "end_time": "2025-04-03T19:35:46.052554", "exception": false, "start_time": "2025-04-03T19:35:46.036348", "status": "completed"}, "tags": []}, "source": ["### ProtoMAML\n", "\n", "A problem of MAML is how to design the output classification layer.\n", "In case all tasks have different number of classes, we need to initialize the output layer with zeros or randomly in every iteration.\n", "Even if we always have the same number of classes, we just start from random predictions.\n", "This requires several inner loop steps to reach a reasonable classification result.\n", "To overcome this problem, Triantafillou et al.\n", "(2020) propose to combine the merits of Prototypical Networks and MAML.\n", "Specifically, we can use prototypes to initialize our output layer to have a strong initialization.\n", "Thereby, it can be shown that the softmax over euclidean distances can be reformulated as a linear layer with softmax.\n", "To see this, let's first write out the negative euclidean distance between a feature vector $f_{\\theta}(\\mathbf{x}^{*})$ of a new data point $\\mathbf{x}^{*}$ to a prototype $\\mathbf{v}_c$ of class $c$:\n", "$$\n", "-||f_{\\theta}(\\mathbf{x}^{*})-\\mathbf{v}_c||^2=-f_{\\theta}(\\mathbf{x}^{*})^Tf_{\\theta}(\\mathbf{x}^{*})+2\\mathbf{v}_c^{T}f_{\\theta}(\\mathbf{x}^{*})-\\mathbf{v}_c^T\\mathbf{v}_c\n", "$$\n", "\n", "We perform the classification across all classes $c\\in\\mathcal{C}$ and take a softmax on the distance.\n", "Hence, any term that is same for all classes can be removed without changing the output probabilities.\n", "In the equation above, this is true for $-f_{\\theta}(\\mathbf{x}^{*})^Tf_{\\theta}(\\mathbf{x}^{*})$ since it is independent of any class prototype.\n", "Thus, we can write:\n", "\n", "$$\n", "-||f_{\\theta}(\\mathbf{x}^{*})-\\mathbf{v}_c||^2=2\\mathbf{v}_c^{T}f_{\\theta}(\\mathbf{x}^{*})-||\\mathbf{v}_c||^2+\\text{constant}\n", "$$\n", "\n", "Taking a second look at the equation above, it looks a lot like a linear layer.\n", "For this, we use $\\mathbf{W}_{c,\\cdot}=2\\mathbf{v}_c$ and $b_c=-||\\mathbf{v}_c||^2$ which gives us the linear layer $\\mathbf{W}f_{\\theta}(\\mathbf{x}^{*})+\\mathbf{b}$.\n", "Hence, if we initialize the output weight with twice the prototypes, and the biases by the negative squared L2 norm of the prototypes, we start with a Prototypical Network.\n", "MAML allows us to adapt this layer and the rest of the network further.\n", "\n", "In the following, we will implement First-Order ProtoMAML for few-shot classification.\n", "The implementation of MAML would be the same except the output layer initialization."]}, {"cell_type": "markdown", "id": "127d61be", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.016006, "end_time": "2025-04-03T19:35:46.084569", "exception": false, "start_time": "2025-04-03T19:35:46.068563", "status": "completed"}, "tags": []}, "source": ["### ProtoMAML implementation\n", "\n", "For implementing ProtoMAML, we can follow Algorithm 2 with minor modifications.\n", "At each training step, we first sample a batch of tasks, and a support and query set for each task.\n", "In our case of few-shot classification, this means that we simply sample multiple support-query set pairs from our sampler.\n", "For each task, we finetune our current model on the support set.\n", "However, since we need to remember the original parameters for the other tasks, the outer loop gradient update and future training steps, we need to create a copy of our model, and finetune only the copy.\n", "We can copy a model by using standard Python functions like `deepcopy`.\n", "The inner loop is implemented in the function `adapt_few_shot` in the PyTorch Lightning module below.\n", "\n", "After finetuning the model, we apply it on the query set and calculate the first-order gradients with respect to the original parameters $\\theta$.\n", "In contrast to simple MAML, we also have to consider the gradients with respect to the output layer initialization, i.e. the prototypes, since they directly rely on $\\theta$.\n", "To realize this efficiently, we take two steps.\n", "First, we calculate the prototypes by applying the original model, i.e. not the copied model, on the support elements.\n", "When initializing the output layer, we detach the prototypes to stop the gradients.\n", "This is because in the inner loop itself, we do not want to consider gradients through the prototypes back to the original model.\n", "However, after the inner loop is finished, we re-attach the computation graph of the prototypes by writing `output_weight = (output_weight - init_weight).detach() + init_weight`.\n", "While this line does not change the value of the variable `output_weight`, it adds its dependency on the prototype initialization `init_weight`.\n", "Thus, if we call `.backward` on `output_weight`, we will automatically calculate the first-order gradients with respect to the prototype initialization in the original model.\n", "\n", "After calculating all gradients and summing them together in the original model, we can take a standard optimizer step.\n", "PyTorch Lightning's method is however designed to return a loss-tensor on which we call `.backward` first.\n", "Since this is not possible here, we need to perform the optimization step ourselves.\n", "All details can be found in the code below.\n", "\n", "For implementing (Proto-)MAML with second-order gradients, it is recommended to use libraries such as [$\\nabla$higher](https://github.com/facebookresearch/higher) from Facebook AI Research.\n", "For simplicity, we stick with first-order methods here."]}, {"cell_type": "code", "execution_count": 25, "id": "6421e405", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:46.118069Z", "iopub.status.busy": "2025-04-03T19:35:46.117619Z", "iopub.status.idle": "2025-04-03T19:35:46.132140Z", "shell.execute_reply": "2025-04-03T19:35:46.131332Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.032782, "end_time": "2025-04-03T19:35:46.133323", "exception": false, "start_time": "2025-04-03T19:35:46.100541", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ProtoMAML(pl.LightningModule):\n", " def __init__(self, proto_dim, lr, lr_inner, lr_output, num_inner_steps):\n", " \"\"\"ProtoMAML.\n", "\n", " Args:\n", " proto_dim: Dimensionality of prototype feature space\n", " lr: Learning rate of the outer loop Adam optimizer\n", " lr_inner: Learning rate of the inner loop SGD optimizer\n", " lr_output: Learning rate for the output layer in the inner loop\n", " num_inner_steps: Number of inner loop updates to perform\n", "\n", " \"\"\"\n", " super().__init__()\n", " self.save_hyperparameters()\n", " self.model = get_convnet(output_size=self.hparams.proto_dim)\n", "\n", " def configure_optimizers(self):\n", " optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr)\n", " scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[140, 180], gamma=0.1)\n", " return [optimizer], [scheduler]\n", "\n", " def run_model(self, local_model, output_weight, output_bias, imgs, labels):\n", " # Execute a model with given output layer weights and inputs\n", " feats = local_model(imgs)\n", " preds = F.linear(feats, output_weight, output_bias)\n", " loss = F.cross_entropy(preds, labels)\n", " acc = (preds.argmax(dim=1) == labels).float()\n", " return loss, preds, acc\n", "\n", " def adapt_few_shot(self, support_imgs, support_targets):\n", " # Determine prototype initialization\n", " support_feats = self.model(support_imgs)\n", " prototypes, classes = ProtoNet.calculate_prototypes(support_feats, support_targets)\n", " support_labels = (classes[None, :] == support_targets[:, None]).long().argmax(dim=-1)\n", " # Create inner-loop model and optimizer\n", " local_model = deepcopy(self.model)\n", " local_model.train()\n", " local_optim = optim.SGD(local_model.parameters(), lr=self.hparams.lr_inner)\n", " local_optim.zero_grad()\n", " # Create output layer weights with prototype-based initialization\n", " init_weight = 2 * prototypes\n", " init_bias = -(torch.norm(prototypes, dim=1) ** 2)\n", " output_weight = init_weight.detach().requires_grad_()\n", " output_bias = init_bias.detach().requires_grad_()\n", "\n", " # Optimize inner loop model on support set\n", " for _ in range(self.hparams.num_inner_steps):\n", " # Determine loss on the support set\n", " loss, _, _ = self.run_model(local_model, output_weight, output_bias, support_imgs, support_labels)\n", " # Calculate gradients and perform inner loop update\n", " loss.backward()\n", " local_optim.step()\n", " # Update output layer via SGD\n", " output_weight.data -= self.hparams.lr_output * output_weight.grad\n", " output_bias.data -= self.hparams.lr_output * output_bias.grad\n", " # Reset gradients\n", " local_optim.zero_grad()\n", " output_weight.grad.fill_(0)\n", " output_bias.grad.fill_(0)\n", "\n", " # Re-attach computation graph of prototypes\n", " output_weight = (output_weight - init_weight).detach() + init_weight\n", " output_bias = (output_bias - init_bias).detach() + init_bias\n", "\n", " return local_model, output_weight, output_bias, classes\n", "\n", " def outer_loop(self, batch, mode=\"train\"):\n", " accuracies = []\n", " losses = []\n", " self.model.zero_grad()\n", "\n", " # Determine gradients for batch of tasks\n", " for task_batch in batch:\n", " imgs, targets = task_batch\n", " support_imgs, query_imgs, support_targets, query_targets = split_batch(imgs, targets)\n", " # Perform inner loop adaptation\n", " local_model, output_weight, output_bias, classes = self.adapt_few_shot(support_imgs, support_targets)\n", " # Determine loss of query set\n", " query_labels = (classes[None, :] == query_targets[:, None]).long().argmax(dim=-1)\n", " loss, preds, acc = self.run_model(local_model, output_weight, output_bias, query_imgs, query_labels)\n", " # Calculate gradients for query set loss\n", " if mode == \"train\":\n", " loss.backward()\n", "\n", " for p_global, p_local in zip(self.model.parameters(), local_model.parameters()):\n", " p_global.grad += p_local.grad # First-order approx. -> add gradients of finetuned and base model\n", "\n", " accuracies.append(acc.mean().detach())\n", " losses.append(loss.detach())\n", "\n", " # Perform update of base model\n", " if mode == \"train\":\n", " opt = self.optimizers()\n", " opt.step()\n", " opt.zero_grad()\n", "\n", " self.log(f\"{mode}_loss\", sum(losses) / len(losses))\n", " self.log(f\"{mode}_acc\", sum(accuracies) / len(accuracies))\n", "\n", " def training_step(self, batch, batch_idx):\n", " self.outer_loop(batch, mode=\"train\")\n", " return None # Returning None means we skip the default training optimizer steps by PyTorch Lightning\n", "\n", " def validation_step(self, batch, batch_idx):\n", " # Validation requires to finetune a model, hence we need to enable gradients\n", " torch.set_grad_enabled(True)\n", " self.outer_loop(batch, mode=\"val\")\n", " torch.set_grad_enabled(False)"]}, {"cell_type": "markdown", "id": "670db1c8", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.01598, "end_time": "2025-04-03T19:35:46.165314", "exception": false, "start_time": "2025-04-03T19:35:46.149334", "status": "completed"}, "tags": []}, "source": ["### Training\n", "\n", "To train ProtoMAML, we need to change our sampling slightly.\n", "Instead of a single support-query set batch, we need to sample multiple.\n", "To implement this, we yet use another Sampler which combines multiple batches from a `FewShotBatchSampler`, and returns it afterwards.\n", "Additionally, we define a `collate_fn` for our data loader which takes the stack of support-query set images, and returns the tasks as a list.\n", "This makes it easier to process in our PyTorch Lightning module before.\n", "The implementation of the sampler can be found below."]}, {"cell_type": "code", "execution_count": 26, "id": "d6ba460b", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:46.199374Z", "iopub.status.busy": "2025-04-03T19:35:46.198659Z", "iopub.status.idle": "2025-04-03T19:35:46.210078Z", "shell.execute_reply": "2025-04-03T19:35:46.209149Z"}, "papermill": {"duration": 0.029543, "end_time": "2025-04-03T19:35:46.211341", "exception": false, "start_time": "2025-04-03T19:35:46.181798", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class TaskBatchSampler:\n", " def __init__(self, dataset_targets, batch_size, N_way, K_shot, include_query=False, shuffle=True):\n", " \"\"\"Task Batch Sampler.\n", "\n", " Args:\n", " dataset_targets: PyTorch tensor of the labels of the data elements.\n", " batch_size: Number of tasks to aggregate in a batch\n", " N_way: Number of classes to sample per batch.\n", " K_shot: Number of examples to sample per class in the batch.\n", " include_query: If True, returns batch of size N_way*K_shot*2, which\n", " can be split into support and query set. Simplifies\n", " the implementation of sampling the same classes but\n", " distinct examples for support and query set.\n", " shuffle: If True, examples and classes are newly shuffled in each\n", " iteration (for training)\n", "\n", " \"\"\"\n", " super().__init__()\n", " self.batch_sampler = FewShotBatchSampler(dataset_targets, N_way, K_shot, include_query, shuffle)\n", " self.task_batch_size = batch_size\n", " self.local_batch_size = self.batch_sampler.batch_size\n", "\n", " def __iter__(self):\n", " # Aggregate multiple batches before returning the indices\n", " batch_list = []\n", " for batch_idx, batch in enumerate(self.batch_sampler):\n", " batch_list.extend(batch)\n", " if (batch_idx + 1) % self.task_batch_size == 0:\n", " yield batch_list\n", " batch_list = []\n", "\n", " def __len__(self):\n", " return len(self.batch_sampler) // self.task_batch_size\n", "\n", " def get_collate_fn(self):\n", " # Returns a collate function that converts one big tensor into a list of task-specific tensors\n", " def collate_fn(item_list):\n", " imgs = torch.stack([img for img, target in item_list], dim=0)\n", " targets = torch.stack([target for img, target in item_list], dim=0)\n", " imgs = imgs.chunk(self.task_batch_size, dim=0)\n", " targets = targets.chunk(self.task_batch_size, dim=0)\n", " return list(zip(imgs, targets))\n", "\n", " return collate_fn"]}, {"cell_type": "markdown", "id": "56e698f3", "metadata": {"papermill": {"duration": 0.015905, "end_time": "2025-04-03T19:35:46.243257", "exception": false, "start_time": "2025-04-03T19:35:46.227352", "status": "completed"}, "tags": []}, "source": ["The creation of the data loaders is with this sampler straight-forward.\n", "Note that since many images need to loaded for a training batch, it is recommended to use less workers than usual."]}, {"cell_type": "code", "execution_count": 27, "id": "116cdebd", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:46.280511Z", "iopub.status.busy": "2025-04-03T19:35:46.279816Z", "iopub.status.idle": "2025-04-03T19:35:46.314907Z", "shell.execute_reply": "2025-04-03T19:35:46.314186Z"}, "papermill": {"duration": 0.054026, "end_time": "2025-04-03T19:35:46.316089", "exception": false, "start_time": "2025-04-03T19:35:46.262063", "status": "completed"}, "tags": []}, "outputs": [], "source": ["# Training constant (same as for ProtoNet)\n", "N_WAY = 5\n", "K_SHOT = 4\n", "\n", "# Training set\n", "train_protomaml_sampler = TaskBatchSampler(\n", " train_set.targets, include_query=True, N_way=N_WAY, K_shot=K_SHOT, batch_size=16\n", ")\n", "train_protomaml_loader = data.DataLoader(\n", " train_set, batch_sampler=train_protomaml_sampler, collate_fn=train_protomaml_sampler.get_collate_fn(), num_workers=2\n", ")\n", "\n", "# Validation set\n", "val_protomaml_sampler = TaskBatchSampler(\n", " val_set.targets,\n", " include_query=True,\n", " N_way=N_WAY,\n", " K_shot=K_SHOT,\n", " batch_size=1, # We do not update the parameters, hence the batch size is irrelevant here\n", " shuffle=False,\n", ")\n", "val_protomaml_loader = data.DataLoader(\n", " val_set, batch_sampler=val_protomaml_sampler, collate_fn=val_protomaml_sampler.get_collate_fn(), num_workers=2\n", ")"]}, {"cell_type": "markdown", "id": "77c3d1e4", "metadata": {"papermill": {"duration": 0.01617, "end_time": "2025-04-03T19:35:46.348522", "exception": false, "start_time": "2025-04-03T19:35:46.332352", "status": "completed"}, "tags": []}, "source": ["Now, we are ready to train our ProtoMAML.\n", "We use the same feature space size as for ProtoNet, but can use a higher learning rate since the outer loop gradients are accumulated over 16 batches.\n", "The inner loop learning rate is set to 0.1, which is much higher than the outer loop lr because we use SGD in the inner loop instead of Adam.\n", "Commonly, the learning rate for the output layer is higher than the base model is the base model is very deep or pre-trained.\n", "However, for our setup, we observed no noticeable impact of using a different learning rate than the base model.\n", "The number of inner loop updates is another crucial hyperparmaeter, and depends on the similarity of our training tasks.\n", "Since all tasks are on images from the same dataset, we notice that a single inner loop update achieves similar performance as 3 or 5 while training considerably faster.\n", "However, especially in RL and NLP, larger number of inner loop steps are often needed."]}, {"cell_type": "code", "execution_count": 28, "id": "ff5d14f3", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:46.383446Z", "iopub.status.busy": "2025-04-03T19:35:46.382593Z", "iopub.status.idle": "2025-04-03T19:35:46.575386Z", "shell.execute_reply": "2025-04-03T19:35:46.574591Z"}, "papermill": {"duration": 0.211289, "end_time": "2025-04-03T19:35:46.576378", "exception": false, "start_time": "2025-04-03T19:35:46.365089", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["GPU available: True (cuda), used: True\n"]}, {"name": "stderr", "output_type": "stream", "text": ["TPU available: False, using: 0 TPU cores\n"]}, {"name": "stderr", "output_type": "stream", "text": ["HPU available: False, using: 0 HPUs\n"]}, {"name": "stderr", "output_type": "stream", "text": ["Lightning automatically upgraded your loaded checkpoint from v1.3.4 to v2.4.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint saved_models/MetaLearning/ProtoMAML.ckpt`\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Found pretrained model at saved_models/MetaLearning/ProtoMAML.ckpt, loading...\n"]}], "source": ["protomaml_model = train_model(\n", " ProtoMAML,\n", " proto_dim=64,\n", " lr=1e-3,\n", " lr_inner=0.1,\n", " lr_output=0.1,\n", " num_inner_steps=1, # Often values between 1 and 10\n", " train_loader=train_protomaml_loader,\n", " val_loader=val_protomaml_loader,\n", ")"]}, {"cell_type": "markdown", "id": "1bde056a", "metadata": {"papermill": {"duration": 0.016519, "end_time": "2025-04-03T19:35:46.609567", "exception": false, "start_time": "2025-04-03T19:35:46.593048", "status": "completed"}, "tags": []}, "source": ["Let's have a look at the training TensorBoard."]}, {"cell_type": "code", "execution_count": 29, "id": "d73b0c4e", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:35:46.643919Z", "iopub.status.busy": "2025-04-03T19:35:46.643515Z", "iopub.status.idle": "2025-04-03T19:35:46.646939Z", "shell.execute_reply": "2025-04-03T19:35:46.646152Z"}, "papermill": {"duration": 0.021959, "end_time": "2025-04-03T19:35:46.648066", "exception": false, "start_time": "2025-04-03T19:35:46.626107", "status": "completed"}, "tags": []}, "outputs": [], "source": ["# Opens tensorboard in notebook. Adjust the path to your CHECKPOINT_PATH if needed\n", "# %tensorboard --logdir ../saved_models/tutorial16/tensorboards/ProtoMAML/"]}, {"cell_type": "markdown", "id": "5b08f321", "metadata": {"papermill": {"duration": 0.016539, "end_time": "2025-04-03T19:35:46.681043", "exception": false, "start_time": "2025-04-03T19:35:46.664504", "status": "completed"}, "tags": []}, "source": ["
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/few-shot-classification.png\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/protonet_classification.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/tensorboard_screenshot_ProtoNet.png\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/MAML_figure.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/MAML_algorithm.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/12-meta-learning/tensorboard_screenshot_ProtoMAML.png\")