{ "cells": [ { "cell_type": "markdown", "id": "586737af", "metadata": {}, "source": [ "# How to use Tune with PyTorch\n", "\n", "(tune-pytorch-cifar-ref)=\n", "\n", "In this walkthrough, we will show you how to integrate Tune into your PyTorch\n", "training workflow. We will follow [this tutorial from the PyTorch documentation](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html)\n", "for training a CIFAR10 image classifier.\n", "\n", "```{image} /images/pytorch_logo.png\n", ":align: center\n", "```\n", "\n", "Hyperparameter tuning can make the difference between an average model and a highly\n", "accurate one. Often simple things like choosing a different learning rate or changing\n", "a network layer size can have a dramatic impact on your model performance. Fortunately,\n", "Tune makes exploring these optimal parameter combinations easy - and works nicely\n", "together with PyTorch.\n", "\n", "As you will see, we only need to add some slight modifications. In particular, we\n", "need to\n", "\n", "1. wrap data loading and training in functions,\n", "2. make some network parameters configurable,\n", "3. add checkpointing (optional),\n", "4. and define the search space for the model tuning\n", "\n", ":::{note}\n", "To run this example, you will need to install the following:\n", "\n", "```bash\n", "$ pip install ray torch torchvision\n", "```\n", ":::\n", "\n", "```{contents}\n", ":backlinks: none\n", ":local: true\n", "```" ] }, { "cell_type": "markdown", "id": "7e8650d1", "metadata": {}, "source": [ "## Setup / Imports\n", "\n", "Let's start with the imports:" ] }, { "cell_type": "code", "execution_count": 1, "id": "55529285", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import os\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import torch.optim as optim\n", "from filelock import FileLock\n", "from torch.utils.data import random_split\n", "import torchvision\n", "import torchvision.transforms as transforms\n", "import ray\n", "from ray import tune\n", "from ray.air import session\n", "from ray.air.checkpoint import Checkpoint\n", "from ray.tune.schedulers import ASHAScheduler" ] }, { "cell_type": "markdown", "id": "f59e551d", "metadata": {}, "source": [ "Most of the imports are needed for building the PyTorch model. Only the last three\n", "imports are for Ray Tune.\n", "\n", "## Data loaders\n", "\n", "We wrap the data loaders in their own function and pass a global data directory.\n", "This way we can share a data directory between different trials." ] }, { "cell_type": "code", "execution_count": 2, "id": "01471556", "metadata": {}, "outputs": [], "source": [ "def load_data(data_dir=\"./data\"):\n", " transform = transforms.Compose([\n", " transforms.ToTensor(),\n", " transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n", " ])\n", "\n", " # We add FileLock here because multiple workers will want to\n", " # download data, and this may cause overwrites since\n", " # DataLoader is not threadsafe.\n", " with FileLock(os.path.expanduser(\"~/.data.lock\")):\n", " trainset = torchvision.datasets.CIFAR10(\n", " root=data_dir, train=True, download=True, transform=transform)\n", "\n", " testset = torchvision.datasets.CIFAR10(\n", " root=data_dir, train=False, download=True, transform=transform)\n", "\n", " return trainset, testset" ] }, { "cell_type": "markdown", "id": "80958cf3", "metadata": {}, "source": [ "## Configurable neural network\n", "\n", "We can only tune those parameters that are configurable. In this example, we can specify\n", "the layer sizes of the fully connected layers:" ] }, { "cell_type": "code", "execution_count": 3, "id": "fff6bd0d", "metadata": {}, "outputs": [], "source": [ "class Net(nn.Module):\n", " def __init__(self, l1=120, l2=84):\n", " super(Net, self).__init__()\n", " self.conv1 = nn.Conv2d(3, 6, 5)\n", " self.pool = nn.MaxPool2d(2, 2)\n", " self.conv2 = nn.Conv2d(6, 16, 5)\n", " self.fc1 = nn.Linear(16 * 5 * 5, l1)\n", " self.fc2 = nn.Linear(l1, l2)\n", " self.fc3 = nn.Linear(l2, 10)\n", "\n", " def forward(self, x):\n", " x = self.pool(F.relu(self.conv1(x)))\n", " x = self.pool(F.relu(self.conv2(x)))\n", " x = x.view(-1, 16 * 5 * 5)\n", " x = F.relu(self.fc1(x))\n", " x = F.relu(self.fc2(x))\n", " x = self.fc3(x)\n", " return x" ] }, { "cell_type": "markdown", "id": "fb619875", "metadata": {}, "source": [ "## The train function\n", "\n", "Now it gets interesting, because we introduce some changes to the example [from the PyTorch\n", "documentation](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html).\n", "\n", "(communicating-with-ray-tune)=\n", "\n", "The full code example looks like this:" ] }, { "cell_type": "code", "execution_count": 4, "id": "fa0bdae0", "metadata": {}, "outputs": [], "source": [ "def train_cifar(config):\n", " net = Net(config[\"l1\"], config[\"l2\"])\n", "\n", " device = \"cpu\"\n", " if torch.cuda.is_available():\n", " device = \"cuda:0\"\n", " if torch.cuda.device_count() > 1:\n", " net = nn.DataParallel(net)\n", " net.to(device)\n", "\n", " criterion = nn.CrossEntropyLoss()\n", " optimizer = optim.SGD(net.parameters(), lr=config[\"lr\"], momentum=0.9)\n", "\n", " # To restore a checkpoint, use `session.get_checkpoint()`.\n", " loaded_checkpoint = session.get_checkpoint()\n", " if loaded_checkpoint:\n", " with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:\n", " model_state, optimizer_state = torch.load(os.path.join(loaded_checkpoint_dir, \"checkpoint.pt\"))\n", " net.load_state_dict(model_state)\n", " optimizer.load_state_dict(optimizer_state)\n", "\n", " data_dir = os.path.abspath(\"./data\")\n", " trainset, testset = load_data(data_dir)\n", "\n", " test_abs = int(len(trainset) * 0.8)\n", " train_subset, val_subset = random_split(\n", " trainset, [test_abs, len(trainset) - test_abs])\n", "\n", " trainloader = torch.utils.data.DataLoader(\n", " train_subset,\n", " batch_size=int(config[\"batch_size\"]),\n", " shuffle=True,\n", " num_workers=8)\n", " valloader = torch.utils.data.DataLoader(\n", " val_subset,\n", " batch_size=int(config[\"batch_size\"]),\n", " shuffle=True,\n", " num_workers=8)\n", "\n", " for epoch in range(10): # loop over the dataset multiple times\n", " running_loss = 0.0\n", " epoch_steps = 0\n", " for i, data in enumerate(trainloader, 0):\n", " # get the inputs; data is a list of [inputs, labels]\n", " inputs, labels = data\n", " inputs, labels = inputs.to(device), labels.to(device)\n", "\n", " # zero the parameter gradients\n", " optimizer.zero_grad()\n", "\n", " # forward + backward + optimize\n", " outputs = net(inputs)\n", " loss = criterion(outputs, labels)\n", " loss.backward()\n", " optimizer.step()\n", "\n", " # print statistics\n", " running_loss += loss.item()\n", " epoch_steps += 1\n", " if i % 2000 == 1999: # print every 2000 mini-batches\n", " print(\"[%d, %5d] loss: %.3f\" % (epoch + 1, i + 1,\n", " running_loss / epoch_steps))\n", " running_loss = 0.0\n", "\n", " # Validation loss\n", " val_loss = 0.0\n", " val_steps = 0\n", " total = 0\n", " correct = 0\n", " for i, data in enumerate(valloader, 0):\n", " with torch.no_grad():\n", " inputs, labels = data\n", " inputs, labels = inputs.to(device), labels.to(device)\n", "\n", " outputs = net(inputs)\n", " _, predicted = torch.max(outputs.data, 1)\n", " total += labels.size(0)\n", " correct += (predicted == labels).sum().item()\n", "\n", " loss = criterion(outputs, labels)\n", " val_loss += loss.cpu().numpy()\n", " val_steps += 1\n", "\n", " # Here we save a checkpoint. It is automatically registered with\n", " # Ray Tune and can be accessed through `session.get_checkpoint()`\n", " # API in future iterations.\n", " os.makedirs(\"my_model\", exist_ok=True)\n", " torch.save(\n", " (net.state_dict(), optimizer.state_dict()), \"my_model/checkpoint.pt\")\n", " checkpoint = Checkpoint.from_directory(\"my_model\")\n", " session.report({\"loss\": (val_loss / val_steps), \"accuracy\": correct / total}, checkpoint=checkpoint)\n", " print(\"Finished Training\")" ] }, { "cell_type": "markdown", "id": "918d8baf", "metadata": {}, "source": [ "As you can see, most of the code is adapted directly from the example.\n", "\n", "## Test set accuracy\n", "\n", "Commonly the performance of a machine learning model is tested on a hold-out test\n", "set with data that has not been used for training the model. We also wrap this in a\n", "function:" ] }, { "cell_type": "code", "execution_count": 5, "id": "93b5b4af", "metadata": {}, "outputs": [], "source": [ "def test_best_model(best_result):\n", " best_trained_model = Net(best_result.config[\"l1\"], best_result.config[\"l2\"])\n", " device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n", " best_trained_model.to(device)\n", "\n", " checkpoint_path = os.path.join(best_result.checkpoint.to_directory(), \"checkpoint.pt\")\n", "\n", " model_state, optimizer_state = torch.load(checkpoint_path)\n", " best_trained_model.load_state_dict(model_state)\n", "\n", " trainset, testset = load_data()\n", "\n", " testloader = torch.utils.data.DataLoader(\n", " testset, batch_size=4, shuffle=False, num_workers=2)\n", "\n", " correct = 0\n", " total = 0\n", " with torch.no_grad():\n", " for data in testloader:\n", " images, labels = data\n", " images, labels = images.to(device), labels.to(device)\n", " outputs = best_trained_model(images)\n", " _, predicted = torch.max(outputs.data, 1)\n", " total += labels.size(0)\n", " correct += (predicted == labels).sum().item()\n", "\n", "\n", " print(\"Best trial test set accuracy: {}\".format(correct / total))" ] }, { "cell_type": "markdown", "id": "85f8230e", "metadata": {}, "source": [ "As you can see, the function also expects a `device` parameter, so we can do the\n", "test set validation on a GPU.\n", "\n", "## Configuring the search space\n", "\n", "Lastly, we need to define Tune's search space. Here is an example:" ] }, { "cell_type": "code", "execution_count": 6, "id": "5416cece", "metadata": {}, "outputs": [], "source": [ "config = {\n", " \"l1\": tune.sample_from(lambda _: 2**np.random.randint(2, 9)),\n", " \"l2\": tune.sample_from(lambda _: 2**np.random.randint(2, 9)),\n", " \"lr\": tune.loguniform(1e-4, 1e-1),\n", " \"batch_size\": tune.choice([2, 4, 8, 16]),\n", "}" ] }, { "cell_type": "markdown", "id": "20af95cc", "metadata": {}, "source": [ "The `tune.sample_from()` function makes it possible to define your own sample\n", "methods to obtain hyperparameters. In this example, the `l1` and `l2` parameters\n", "should be powers of 2 between 4 and 256, so either 4, 8, 16, 32, 64, 128, or 256.\n", "The `lr` (learning rate) should be uniformly sampled between 0.0001 and 0.1. Lastly,\n", "the batch size is a choice between 2, 4, 8, and 16.\n", "\n", "At each trial, Tune will now randomly sample a combination of parameters from these\n", "search spaces. It will then train a number of models in parallel and find the best\n", "performing one among these. We also use the `ASHAScheduler` which will terminate bad\n", "performing trials early.\n", "\n", "You can specify the number of CPUs, which are then available e.g.\n", "to increase the `num_workers` of the PyTorch `DataLoader` instances. The selected\n", "number of GPUs are made visible to PyTorch in each trial. Trials do not have access to\n", "GPUs that haven't been requested for them - so you don't have to care about two trials\n", "using the same set of resources.\n", "\n", "Here we can also specify fractional GPUs, so something like `gpus_per_trial=0.5` is\n", "completely valid. The trials will then share GPUs among each other.\n", "You just have to make sure that the models still fit in the GPU memory.\n", "\n", "After training the models, we will find the best performing one and load the trained\n", "network from the checkpoint file. We then obtain the test set accuracy and report\n", "everything by printing.\n", "\n", "The full main function looks like this:" ] }, { "cell_type": "code", "execution_count": 7, "id": "91d83380", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2022-07-22 16:38:53,384\tINFO services.py:1483 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8273\u001b[39m\u001b[22m\n", "2022-07-22 16:38:56,785\tWARNING function_trainable.py:619 -- Function checkpointing is disabled. This may result in unexpected behavior when using checkpointing features or certain schedulers. To enable, set the train function arguments to be `func(config, checkpoint_dir=None)`.\n" ] }, { "data": { "text/html": [ "== Status ==
Current time: 2022-07-22 16:40:13 (running for 00:01:16.43)
Memory usage on this node: 10.7/16.0 GiB
Using AsyncHyperBand: num_stopped=2\n", "Bracket: Iter 2.000: -1.421571186053753 | Iter 1.000: -1.7652838359832763
Resources requested: 0/16 CPUs, 0/0 GPUs, 0.0/5.63 GiB heap, 0.0/2.0 GiB objects
Current best trial: 66098_00000 with loss=1.421571186053753 and parameters={'l1': 128, 'l2': 128, 'lr': 0.00046907397024184945, 'batch_size': 4}
Result logdir: /Users/kai/ray_results/train_cifar_2022-07-22_16-38-50
Number of trials: 2/2 (2 TERMINATED)
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
Trial name status loc batch_size l1 l2 lr iter total time (s) loss accuracy
train_cifar_66098_00000TERMINATED127.0.0.1:53065 4 128 1280.000469074 2 72.61761.42157 0.4877
train_cifar_66098_00001TERMINATED127.0.0.1:53078 4 128 640.00993903 1 64.97211.90462 0.2915


" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "2022-07-22 16:38:57,794\tINFO plugin_schema_manager.py:52 -- Loading the default runtime env schemas: ['/Users/kai/coding/ray/python/ray/_private/runtime_env/../../runtime_env/schemas/working_dir_schema.json', '/Users/kai/coding/ray/python/ray/_private/runtime_env/../../runtime_env/schemas/pip_schema.json'].\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\u001b[2m\u001b[36m(train_cifar pid=53065)\u001b[0m Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to /Users/kai/ray_results/train_cifar_2022-07-22_16-38-50/train_cifar_66098_00000_0_batch_size=4,l1=128,l2=128,lr=0.0005_2022-07-22_16-38-57/data/cifar-10-python.tar.gz\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/170498071 [00:00