{ "cells": [ { "cell_type": "markdown", "id": "ebf9c0f5", "metadata": {}, "source": [ "# Implementing a CNN\n", "\n", "Now, we shall undertake implementing a CNN for image classification.\n", "\n", "1. Although `pytorch` has built-in convolution and pooling layers, those don't support ensembles, so we'll have to make our own.\n", "2. I give the indices relevant to convolution in the lecture notes. But indexing with these would result in copying tensors, thus wasting memory and compute time. However, the built-in convolution and pooling operations do are effective, so we can use those." ] }, { "cell_type": "markdown", "id": "bf411eed", "metadata": {}, "source": [ "## Setup\n", "\n", "### Imports\n", "\n", "Import `defaultdict`, `Callable`, `datasets`, `math`, `torch`, `torch.nn.functional` as `F`, `tqdm` and `Optional`.\n", "\n", "Moreover, import:\n", "1. The functions `get_accuracy`, `get_cross_entropy`, and `get_dataloader_random_reshuffle`, that you wrote in Notebook 0623,\n", "1. The class `Linear`, that you wrote in Notebook 0706.\n", "1. The functions `evaluate_model`, `normalize_features`, `pbt_init` and `pbt_update`, that you wrote in Notebook 0707.\n", "1. The classes `AdamW` and `Optimizer`, that you wrote in Notebook 0708." ] }, { "cell_type": "code", "execution_count": null, "id": "d0dadd32", "metadata": {}, "outputs": [], "source": [ "from collections import defaultdict\n", "from collections.abc import (\n", " Callable\n", ")\n", "import datasets\n", "import math\n", "import torch\n", "import torch.nn.functional as F\n", "import tqdm\n", "from typing import Optional\n", "\n", "from util_0713 import (\n", " AdamW,\n", " evaluate_model,\n", " get_accuracy,\n", " get_cross_entropy,\n", " get_dataloader_random_reshuffle,\n", " Linear,\n", " normalize_features,\n", " Optimizer,\n", " pbt_init,\n", " pbt_update,\n", ")" ] }, { "cell_type": "markdown", "id": "b72a3041", "metadata": {}, "source": [ "### Configuration\n", "\n", "Create a configuration dictionary with the following keys:\n", "- `\"dataset_path\"`: `str` \n", " As we achieved 98% accuracy on MNIST with an MLP and AdamW in Notebook 0708, we'll take a step up and train on CIFAR-10: https://huggingface.co/datasets/uoft-cs/cifar10 [10], that we used previously in Homework 2.\n", "- `\"device\"`: `torch.device | int | str` \n", " The device identifier.\n", "- `\"ensemble_shape\"`: `tuple[int]` \n", " Make this `(16,)`.\n", "- `\"hyperparameter_raw_init_distributions\"`, `\"hyperparameter_raw_perturb\"`, `\"hyperparameter_transforms\"` : `dict` \n", " These three dictionaries are going to determine how the hyperparameters are tuned. We'll tune the following hyperparameters:\n", " 1. Epsilon $\\epsilon$.\n", " 2. Learning rate $\\eta$.\n", " 3. Weight decay $\\lambda$.\n", " 4. First moment moving average decay rate $\\beta_1$.\n", " 5. Second moment moving average decay rate $\\beta_2$.\n", "\n", " Of these, we don't know the required order of magnitude of the first three. Thus it may be good to make them distributed along $10^\\mathscr D$ where $\\mathscr D$ is a normal or uniform distribution. You can try to center the distributions at the recommended values.\n", "\n", " We know that the recommended values of the fourth and fifth are $0.9$ and $0.999$. So it may be best to give them a distribution of the form $1-10^\\mathscr D$.\n", "- `\"improvement_threshold:`: `float` \n", " Make this `1e-4`.\n", "- `\"minibatch_size\"`: `int` \n", " Make this a `64`.\n", "- `\"minibatch_size_eval\"`: `int` \n", " On my home computer, I can make this `128`.\n", "- `\"pbt\"` : `bool` \n", " Make this `True`.\n", "- `\"seed\"`: `int` \n", " This is for reproducible experiments. Insert any integer.\n", "- `\"steps_num\"`: `int` \n", " To reach at least 80% accuracy, you probably need about `100_001` steps, which take about an hour on GPU. Just to see learning take effect, you can make this `10_001`.\n", "- `\"steps_without_improvement`: `int` \n", " Make this `10_000`.\n", "- `\"valid_interval\"`: `int` \n", " Make this `1000`.\n", "- `\"welch_confidence_level\"`: `float` \n", " We will exploit based on a one-sided Welch $t$-test with this confidence level. Maybe you can try `.8`. Feel free to try out various values here!\n", "- `\"welch_sample_size\"`: `int` \n", " We will exploit based on a one-sided Welch $t$-test on the last this many validation metrics of the population members. To follow the PBT paper, make this `10`." ] }, { "cell_type": "code", "execution_count": null, "id": "47fa2a3c", "metadata": {}, "outputs": [], "source": [ "config = {\n", " \"dataset_path\": \"uoft-cs/cifar10\",\n", " \"device\": \"cuda\",\n", " \"ensemble_shape\": (16,),\n", " \"hyperparameter_raw_init_distributions\": {\n", " \"epsilon\": torch.distributions.Uniform(\n", " torch.tensor(-10, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(-5, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"first_moment_decay\": torch.distributions.Uniform(\n", " torch.tensor(-3, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"learning_rate\": torch.distributions.Uniform(\n", " torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"second_moment_decay\": torch.distributions.Uniform(\n", " torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"weight_decay\": torch.distributions.Uniform(\n", " torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n", " )\n", " },\n", " \"hyperparameter_raw_perturb\": {\n", " \"epsilon\": torch.distributions.Normal(\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"first_moment_decay\": torch.distributions.Normal(\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"learning_rate\": torch.distributions.Normal(\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"second_moment_decay\": torch.distributions.Normal(\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " \"weight_decay\": torch.distributions.Normal(\n", " torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n", " torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n", " ),\n", " },\n", " \"hyperparameter_transforms\": {\n", " \"epsilon\": lambda log10: 10 ** log10,\n", " \"first_moment_decay\": lambda x: (1 - 10 ** x).clamp(0, 1),\n", " \"learning_rate\": lambda log10: 10 ** log10,\n", " \"second_moment_decay\": lambda x: (1 - 10 ** x).clamp(0, 1),\n", " \"weight_decay\": lambda log10: 10 ** log10,\n", " },\n", " \"improvement_threshold\": 1e-4,\n", " \"minibatch_size\": 64,\n", " \"minibatch_size_eval\": 1 << 7,\n", " \"pbt\": True,\n", " \"seed\": 0,\n", " \"steps_num\": 10_001,\n", " \"steps_without_improvement\": 10_000,\n", " \"valid_interval\": 1000,\n", " \"welch_confidence_level\": .8,\n", " \"welch_sample_size\": 10,\n", "}" ] }, { "cell_type": "markdown", "id": "d0974b31", "metadata": {}, "source": [ "Set the `torch` pseudo-random number generation seed as per the configuration dictionary." ] }, { "cell_type": "code", "execution_count": null, "id": "37fe9a06", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(config[\"seed\"])" ] }, { "cell_type": "markdown", "id": "7d15ed55", "metadata": {}, "source": [ "## Dataset\n", "\n", "### Load the Dataset\n", "\n", "Load the dataset specified by the configuration dictionary. Make a train-valid split of its `\"train\"` split. Print tensor shapes and datatypes." ] }, { "cell_type": "code", "execution_count": null, "id": "35775427", "metadata": {}, "outputs": [], "source": [ "dataset = datasets.load_dataset(\n", " config[\"dataset_path\"]\n", ").with_format(\n", " \"torch\",\n", " device=config[\"device\"]\n", ")\n", "print(dataset)" ] }, { "cell_type": "code", "execution_count": null, "id": "50e8e0d3", "metadata": {}, "outputs": [], "source": [ "train = dataset[\"train\"]\n", "train_valid = train.train_test_split(\n", " seed=config[\"seed\"],\n", " test_size=10_000,\n", ")\n", "train, valid = (\n", " train_valid[key][:]\n", " for key in [\"train\", \"test\"]\n", ")\n", "for split in (train, valid):\n", " for key in split.keys():\n", " entries = split[key]\n", " print(key, entries.shape, entries.dtype)" ] }, { "cell_type": "markdown", "id": "f8884c41", "metadata": {}, "source": [ "Note that the shapes of the images are `(batch dimension, channel dimension, vertical dimension, horizontal dimension)`. Usually, I prefer feature dimensions (here: channel) after spatial dimensions. But this time, to be more compatible with `torch` conventions, we'll keep the original order.\n", "\n", "Still, convert the feature tensors to floating points and normalize them." ] }, { "cell_type": "code", "execution_count": null, "id": "ede63da2", "metadata": {}, "outputs": [], "source": [ "train_features, valid_features = (\n", " d[\"img\"].to(torch.float32)\n", " for d in (train, valid)\n", ")\n", "normalize_features(\n", " train_features,\n", " (valid_features,),\n", " verbose=True\n", ")\n", "print(train_features.mean(), train_features.std())" ] }, { "cell_type": "markdown", "id": "0de04aab", "metadata": {}, "source": [ "Create a training dataloader and print minibatch entry shapes, to see if the machinery is working as expected." ] }, { "cell_type": "code", "execution_count": null, "id": "e8bfe1c6", "metadata": {}, "outputs": [], "source": [ "train_labels = train[\"label\"]\n", "valid_labels = valid[\"label\"]\n", "\n", "dataloader_train = get_dataloader_random_reshuffle(\n", " config,\n", " train_features,\n", " train_labels\n", ")\n", "minibatch_features, minibatch_labels = next(dataloader_train)\n", "print(minibatch_features.shape)\n", "print(minibatch_labels.shape)" ] }, { "cell_type": "markdown", "id": "8fd9998b", "metadata": {}, "source": [ "## New Layers\n", "\n", "Let's implement the new, 2d convolution and 2d mean pool layers." ] }, { "cell_type": "markdown", "id": "d2cc5d0c", "metadata": {}, "source": [ "### `Conv2D`\n", "\n", "Let's write our own convolution layer! As we'll want to train an ensemble, we can't use the built-in module `torch.nn.Module.Conv2D`. However, we can make do with the function `F.conv2d` as it supports blocked connections via the `groups` keyword argument. [Follow this link](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv2d.html) for the documentation. You can implement the ensemble-ready two-dimensional convolution layer using this function and some tensor shape transformations.\n", "\n", "Make the layer work both with ensembled and non-ensembled input.\n", "\n", "When you're finished, make a 2D convolution layer with kernel shape `(3,3)` and check the shape of its output features on a training minibatch, and a slice of the validation features." ] }, { "cell_type": "code", "execution_count": null, "id": "43b55e7c", "metadata": {}, "outputs": [], "source": [ "class Conv2D(torch.nn.Module):\n", " \"\"\"\n", " Ensemble-ready, two-dimensional convolution layer\n", "\n", " Arguments\n", " ---------\n", " config : `dict`\n", " Configuration dictionary. Required key-value pairs:\n", " `\"device\"` : `str`\n", " The device to store parameters on.\n", " `\"ensemble_shape\"` : `tuple[int]`\n", " The shape of the ensemble of affine transformations\n", " the model represents.\n", " in_channels : `int`\n", " The number of input channels.\n", " kernel_shape : `tuple[int]`\n", " The kernel shape.\n", " out_channels : `int`\n", " The number of output channels.\n", " bias : `bool`, optional\n", " Whether to include bias along the output channels.\n", " dilation : `int | tuple[int]`, optional\n", " The spacing between kernel elements, in all directions,\n", " or per direction. Default: `1`.\n", " init_multiplier : `float`, optional\n", " We initialize linear maps with Glorot normal initialization,\n", " that is using the centered normal distribution\n", " with standard deviation `out_channels ** -.5` times this value.\n", " Default: `1.`.\n", " padding : `int | str | tuple[int]`, optional\n", " The stride in all directions or per direction.\n", " Alternatively, `\"valid\"` is the same as `0`,\n", " and `\"same\"` pads the input so the output has the same shape\n", " as the input.\n", " Default: `0`.\n", " stride : `int | tuple[int]`, optional\n", " The stride in all directions or per direction.\n", " Default: `1`.\n", "\n", " Calling\n", " -------\n", " Instance calls require one positional argument:\n", " features : `torch.Tensor`\n", " Tensor of features, of shape\n", " `batch_shape + (in_channels, height, width)` or\n", " `ensemble_shape + batch_shape + (in_channels, height, width)`\n", " \"\"\"\n", " def __init__(\n", " self,\n", " config: dict,\n", " in_channels: int,\n", " kernel_shape: tuple[int],\n", " out_channels: int,\n", " bias=True,\n", " dilation=1,\n", " init_multiplier=1.,\n", " padding=0,\n", " stride=1\n", " ):\n", " super().__init__()\n", "\n", " self.dilation = dilation\n", " self.ensemble_shape = config[\"ensemble_shape\"]\n", " self.in_channels = in_channels\n", " self.kernel_shape = kernel_shape\n", " self.out_channels = out_channels\n", " self.padding = padding\n", " self.stride = stride\n", "\n", " height, width = kernel_shape\n", " self.weight = torch.nn.Parameter(torch.empty(\n", " self.ensemble_shape\n", " + (\n", " out_channels,\n", " in_channels,\n", " height,\n", " width\n", " ),\n", " device=config[\"device\"],\n", " dtype=torch.float32\n", " ).normal_(std=out_channels ** -.5) * init_multiplier)\n", "\n", " if bias:\n", " self.bias = torch.nn.Parameter(torch.zeros(\n", " self.ensemble_shape + (out_channels,),\n", " device=config[\"device\"],\n", " dtype=torch.float32\n", " ))\n", " else:\n", " self.bias = None\n", "\n", "\n", " def forward(self, features: torch.Tensor) -> torch.Tensor:\n", " ensemble_dim = len(self.ensemble_shape)\n", "\n", " if features.ndim == 4:\n", " features = features.expand(\n", " self.ensemble_shape + features.shape\n", " )\n", "\n", " features = (\n", " features\n", " .movedim(ensemble_dim, 0)\n", " .flatten(1, ensemble_dim + 1)\n", " ) # (minibatch_size, in_channels_total, input_height, input_width)\n", "\n", " features = F.conv2d(\n", " features,\n", " self.weight.flatten(end_dim=ensemble_dim),\n", " bias=self.bias.flatten(end_dim=ensemble_dim),\n", " dilation=self.dilation,\n", " groups=max(1, sum(self.ensemble_shape)),\n", " padding=self.padding,\n", " stride=self.stride\n", " ) # (minibatch_size, out_channels_total, output_height, output_width)\n", "\n", " features = (\n", " features\n", " .unflatten(1, self.ensemble_shape + (self.out_channels,))\n", " .movedim(0, ensemble_dim)\n", " )\n", "\n", " return features\n", " \n", "layer = Conv2D(\n", " config,\n", " 3,\n", " (3,3),\n", " 16\n", ")\n", "print(layer(minibatch_features).shape)\n", "print(layer(valid_features[0:config[\"minibatch_size_eval\"]]).shape)" ] }, { "cell_type": "markdown", "id": "3990dc7f", "metadata": {}, "source": [ "### `Pool2D`\n", "\n", "Let's now write a mean pool operation! It can be created similarly to `Conv2D`, this time using `F.avg_pool2d` in case we are pooling along the kernel displacements.\n", "\n", "When finished test output shapes giving either\n", "1. `kernel_shape` or\n", "2. `sequence_dim_num`." ] }, { "cell_type": "code", "execution_count": null, "id": "b0669ed9", "metadata": {}, "outputs": [], "source": [ "class Pool2D(torch.nn.Module):\n", " \"\"\"\n", " Ensemble-ready two-dimensional mean pool operation\n", "\n", " Arguments\n", " ---------\n", " config : `dict`\n", " Configuration dictionary. Required key-value pairs:\n", " `\"device\"` : `str`\n", " The device to store parameters on.\n", " `\"ensemble_shape\"` : `tuple[int]`\n", " The shape of the ensemble of affine transformations\n", " the model represents.\n", " kernel_shape : `int | tuple[int]`, optional\n", " The kernel shape.\n", " If given, we pool along the kernel displacements.\n", " Otherwise, we pool along all the two sequential dimensions.\n", " padding : `int | tuple[int]`, optional\n", " The padding in all directions or per direction.\n", " It is used if `kernel_shape` is given.\n", " Default: `0`.\n", " stride : `int | tuple[int]`, optional\n", " The stride in all directions or per direction.\n", " It is used if `kernel_shape` is given.\n", " Default: `1`.\n", "\n", " Calling\n", " -------\n", " Instance calls require one positional argument:\n", " features : `torch.Tensor`\n", " Tensor of features, of shape\n", " `batch_shape + (in_channels, height, width)` or\n", " `ensemble_shape + batch_shape + (in_channels, height, width)`\n", " \"\"\"\n", " def __init__(\n", " self,\n", " config: dict,\n", " kernel_shape: Optional[tuple[int]] = None,\n", " padding=0,\n", " stride=1\n", " ):\n", " super().__init__()\n", "\n", " self.ensemble_shape = config[\"ensemble_shape\"]\n", " self.kernel_shape = kernel_shape\n", " self.padding = padding\n", " self.stride = stride\n", "\n", "\n", " def forward(self, features: torch.Tensor) -> torch.Tensor:\n", " if features.ndim == 4:\n", " features = features.expand(\n", " self.ensemble_shape + features.shape\n", " )\n", " \n", " if self.kernel_shape is None:\n", " return features.mean(dim=(-2, -1))\n", "\n", " channels = features.shape[-3]\n", " ensemble_dim = len(self.ensemble_shape)\n", " features = (\n", " features\n", " .movedim(ensemble_dim, 0)\n", " .flatten(1, ensemble_dim + 1)\n", " ) # (minibatch_size, channels_total, input_height, input_width)\n", " features = F.avg_pool2d(\n", " features,\n", " self.kernel_shape,\n", " padding=self.padding,\n", " stride=self.stride\n", " ) # (minibatch_size, channels_total, output_height, output_width)\n", " features = (\n", " features\n", " .unflatten(1, self.ensemble_shape + (channels,))\n", " .movedim(0, ensemble_dim)\n", " )\n", "\n", " return features\n", " \n", "\n", "layer = Pool2D(config, kernel_shape=(3,3), stride=2)\n", "print(layer(minibatch_features).shape)\n", "print(layer(valid_features[0:config[\"minibatch_size_eval\"]]).shape)\n", "\n", "layer = Pool2D(config)\n", "print(layer(minibatch_features).shape)\n", "print(layer(valid_features[0:config[\"minibatch_size_eval\"]]).shape)" ] }, { "cell_type": "markdown", "id": "b88fc528", "metadata": {}, "source": [ "### Make a CNN\n", "\n", "Time to make a full CNN! There's a lot of create freedom here. The general idea is to alternate the following:\n", "1. Process local data by one or more `Conv2D` layers, with ReLUs in between.\n", "2. Aggregate data by a `Pool2D` layer of stride > 1.\n", "\n", "Rules of thumb:\n", "1. As you aggregate data you can have more channels.\n", "2. Keep kernel sizes small odd numbers larger than 1.\n", "3. The last `Pool2D` layer should be a total pooling one, with `kernel_shape` not given. It should be followed by a small MLP.\n", "\n", "Make a model, then check the shape of its output on a minibatch. The output should have shape appropriate for logits." ] }, { "cell_type": "code", "execution_count": null, "id": "dd763da9", "metadata": {}, "outputs": [], "source": [ "model = torch.nn.Sequential(\n", " Conv2D(\n", " config,\n", " 3,\n", " (3,3),\n", " 16,\n", " init_multiplier=2 ** .5\n", " ),\n", " torch.nn.ReLU(),\n", " Conv2D(\n", " config,\n", " 16,\n", " (3,3),\n", " 32,\n", " # init_multiplier=2 ** .5\n", " ),\n", " Pool2D(\n", " config,\n", " kernel_shape=(3,3),\n", " stride=2\n", " ),\n", " Conv2D(\n", " config,\n", " 32,\n", " (3,3),\n", " 64,\n", " init_multiplier=2 ** .5\n", " ),\n", " torch.nn.ReLU(),\n", " Conv2D(\n", " config,\n", " 64,\n", " (3,3),\n", " 128\n", " ),\n", " Pool2D(\n", " config\n", " ),\n", " Linear(\n", " config,\n", " 128,\n", " 128,\n", " init_multiplier=2 ** .5\n", " ),\n", " torch.nn.ReLU(),\n", " Linear(\n", " config,\n", " 128,\n", " 10\n", " )\n", ")\n", "\n", "print(model(minibatch_features).shape)" ] }, { "cell_type": "markdown", "id": "26652068", "metadata": {}, "source": [ "## Train a CNN\n", "\n", "Time to train this thing! You can start out from the version of the function `pbt` in Notebook 1105. Make the following changes:\n", "1. Some hyperparameter configurations can make training numerically unstable, which results in NaN (Not a Number) values. To avoid propagation of these values, we can convert them to the worst possible losses and metrics in PBT updates, so that population members with NaN values get replaced by those without. To this end, when you calculate training and validation losses and metrics in the beginning of the evaluation phase of the function, apply the `nan_to_num` method of the loss and metric tensors to convert NaN values to $\\infty$ and $-\\infty$, respectively.\n", "2. Where you check if one of the best last metrics is better than the all-time best, you should also record the population member ID that has the best last metric. If the best last metric does is better than the all-time best, then when you update the all-time best, also record at the `\"best parameters\"` key of the log dictionary the parameters of the best model. You can do this by looping over the key-value pairs of the `state_dict` of the model, and taking the subtensors of the parameters given by the best metric population ID.\n", "\n", "Create an `AdamW` optimizer for the model parameters and run training." ] }, { "cell_type": "code", "execution_count": null, "id": "ff8ebcc3", "metadata": {}, "outputs": [], "source": [ "def train_supervised(\n", " config: dict,\n", " get_loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],\n", " get_metric: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],\n", " model: torch.nn.Module,\n", " optimizer: Optimizer,\n", " train_features: torch.Tensor,\n", " train_values: torch.Tensor,\n", " valid_features: torch.Tensor,\n", " valid_values: torch.Tensor,\n", ") -> dict:\n", " \"\"\"\n", " Population-based training on a supervised learning task.\n", " Tuned hyperparameters are given by raw values and transformations.\n", " This way, the hyperparameters are perturbed by\n", " additive noise on raw values.\n", "\n", " Parameters\n", " ----------\n", " config : `dict`\n", " Configuration dictionary. Required key-value pairs:\n", " `\"ensemble_shape\"` : tuple[int]\n", " Ensemble shape. We assume this is a 1-dimensional tuple\n", " with dimensions the population size.\n", " `\"hyperparameter_raw_init_distributions\"` : `dict`\n", " Dictionary that maps tuned hyperparameter names\n", " to `torch.distributions.Distribution` of raw hyperparameter values.\n", " Required keys:\n", " `\"learning_rate\"`:\n", " The learning rate of stochastic gradient descent.\n", " `\"hyperparameter_raw_perturbs\"` : `dict`\n", " Dictionary that maps tuned hyperparameter names\n", " to `torch.distributions.Distribution` of additive noise.\n", " `\"hyperparameter_transforms\"` : `dict`\n", " Dictionary that maps tuned hyperparameter names\n", " to transformations of raw hyperparameter values.\n", " `\"improvement_threshold\"` : `float`\n", " A new metric score has to be this much better\n", " than the previous best to count as an improvement.\n", " `\"minibatch_size\"` : `int`\n", " Minibatch size to use in a training step.\n", " `\"minibatch_size_eval\"` : `int`\n", " Minibatch size to use in evaluation.\n", " On CPU, should be about the same as `minibatch_size`.\n", " On GPU, should be as big as possible without\n", " incurring an Out of Memory error.\n", " `\"pbt\"` : `bool`\n", " Whether to use PBT updates in validations.\n", " If `False`, the algorithm just samples hyperparameters at start,\n", " then keeps them constant.\n", " `\"steps_num\"` : `int`\n", " Maximum number of training steps.\n", " `\"steps_without_improvement`\" : `int`\n", " If the number of training steps without improvement\n", " exceeds this value, then training is stopped.\n", " `\"valid_interval\"` : `int`\n", " Frequency of evaluations, measured in number of training steps.\n", " `\"welch_confidence_level\"` : `float`\n", " The confidence level in Welch's t-test\n", " that is used in determining if a population member\n", " is to be replaced by another member with perturbed hyperparameters.\n", " `\"welch_sample_size\"` : `int`\n", " The last this many validation metrics are used\n", " in Welch's t-test.\n", " `get_loss` : `Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`\n", " A function that maps a pair of predicted and target value tensors\n", " to a tensor of losses per ensemble member.\n", " `get_metric` : `Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`\n", " A function that maps a pair of predicted and target value tensors\n", " to a tensor of metrics per ensemble member.\n", " We assume a greater metric is better.\n", " `model` : `torch.nn.Module`\n", " The model ensemble to tune.\n", " `optimizer` : `Optimizer`\n", " An optimizer that tracks the parameters of `model`.\n", " `train_features` : `torch.Tensor`\n", " Training feature tensor.\n", " `train_values` : `torch.Tensor`\n", " Training value tensor.\n", " `valid_features` : `torch.Tensor`\n", " Validation feature tensor.\n", " `valid_values` : `torch.Tensor`\n", " Validation value tensor.\n", " \n", " Returns\n", " -------\n", " An output dictionary with the following key-value pairs:\n", " `\"best parameters\"` : `dict` \n", " The state dictionary of the model with the best metric\n", " encountered during training.\n", " `\"source mask\"` : `torch.Tensor`\n", " The source masks of population members\n", " that were replace by other members in a PBT update\n", " `\"target indices\"` : `torch.Tensor`\n", " The indices of population members\n", " that the member where the source mask is to were replaced with.\n", " `\"validation metric\"` : `torch.Tensor`\n", " The validation metrics at evaluation steps.\n", "\n", " In addition, for each tuned hyperparameter name,\n", " we include a `torch.Tensor` of values per update.\n", " \"\"\"\n", " ensemble_shape = config[\"ensemble_shape\"]\n", " if len(ensemble_shape) != 1:\n", " raise ValueError(f\"The number of dimensions in the ensemble shape should be 1 for the population size, but it is {len(ensemble_shape)}\")\n", "\n", " population_size = ensemble_shape[0]\n", " config_local = dict(config)\n", " log = defaultdict(list)\n", "\n", " pbt_init(config_local, log)\n", "\n", " optimizer.update_config(config_local)\n", "\n", " best_valid_metric = -torch.inf\n", " progress_bar = tqdm.trange(config[\"steps_num\"])\n", " steps_without_improvement = 0\n", " train_dataloader = get_dataloader_random_reshuffle(\n", " config,\n", " train_features,\n", " train_values\n", " )\n", "\n", " for step_id in progress_bar: \n", " if step_id % config[\"valid_interval\"] == 0:\n", " with torch.no_grad():\n", " validation_metric = evaluate_model(\n", " config,\n", " valid_features,\n", " get_metric,\n", " model,\n", " valid_values\n", " ).nan_to_num(-torch.inf)\n", " log[\"validation metric\"].append(validation_metric)\n", " print(\n", " f\"validation metric {validation_metric.max().cpu().item():.4f}\"\n", " )\n", "\n", " best_last_metric, best_last_metric_id \\\n", " = log[\"validation metric\"][-1].max(dim=-1)\n", " print(\n", " f\"Best last metric {best_last_metric.cpu().item():.2f}\",\n", " flush=True\n", " )\n", " if (\n", " best_valid_metric + config[\"improvement_threshold\"]\n", " ) < best_last_metric:\n", " print(\n", " f\"New best metric\",\n", " flush=True\n", " )\n", " best_valid_metric = best_last_metric\n", " steps_without_improvement = 0\n", " log[\"best parameters\"] = {\n", " key: value[best_last_metric_id].clone()\n", " for key, value in model.state_dict().items()\n", " }\n", " else:\n", " print(\n", " f\"Best metric {best_valid_metric.cpu().item():.2f}\",\n", " flush=True\n", " )\n", " steps_without_improvement += config[\"valid_interval\"]\n", " if steps_without_improvement > config[\n", " \"steps_without_improvement\"\n", " ]:\n", " break\n", "\n", " if config[\"pbt\"] and (len(log[\"validation metric\"]) >= config[\n", " \"welch_sample_size\"\n", " ]):\n", " evaluations = torch.stack(\n", " log[\"validation metric\"][-config[\"welch_sample_size\"]:]\n", " )\n", " pbt_update(\n", " config_local, evaluations, log, optimizer.get_parameters()\n", " )\n", "\n", " optimizer.update_config(config_local)\n", "\n", " minibatch_features, minibatch_labels = next(train_dataloader)\n", " optimizer.zero_grad()\n", "\n", " predict = model(minibatch_features)\n", "\n", " loss = get_loss(predict, minibatch_labels).sum()\n", " loss.backward()\n", " optimizer.step()\n", "\n", "\n", " progress_bar.close()\n", " for key, value in log.items():\n", " if isinstance(value, list):\n", " log[key] = torch.stack(value)\n", "\n", " return log\n", "\n", "optimizer = AdamW(model.parameters())\n", "log = train_supervised(\n", " config,\n", " get_cross_entropy,\n", " get_accuracy,\n", " model,\n", " optimizer,\n", " train_features,\n", " train_labels,\n", " valid_features,\n", " valid_labels\n", ")" ] }, { "cell_type": "markdown", "id": "97036bf4", "metadata": {}, "source": [ "## Dataset Reference\n", "\n", "[1] Alex Krizhevsky: *Learning Multiple Layers of Features from Tiny Images*. 2009. https://www.cs.toronto.edu/~kriz/cifar.html" ] }, { "cell_type": "markdown", "id": "a1cf38c4", "metadata": {}, "source": [ "## License\n", "\n", "This work is licensed under Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International. To view a copy of this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/" ] }, { "cell_type": "code", "execution_count": null, "id": "bea5abbb", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "dml", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.14" } }, "nbformat": 4, "nbformat_minor": 5 }