{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "c6962854",
            "metadata": {},
            "source": [
                "# Training a Torch Image Classifier\n",
                "\n",
                "This tutorial shows you how to train an image classifier using the [Ray AI Runtime](air) (AIR).\n",
                "\n",
                "You should be familiar with [PyTorch](https://pytorch.org/) before starting the tutorial. If you need a refresher, read PyTorch's [training a classifier](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html) tutorial.\n",
                "\n",
                "## Before you begin\n",
                "\n",
                "* Install the [Ray AI Runtime](air). You need Ray 2.0 or later to run this example."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "d806ba6b",
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install 'ray[air]'"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "6d588ce2",
            "metadata": {},
            "source": [
                "* Install `requests`, `torch`, and `torchvision`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "77a70a7a",
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install requests torch torchvision"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f18ec14f",
            "metadata": {},
            "source": [
                "## Load and normalize CIFAR-10\n",
                "\n",
                "We'll train our classifier on a popular image dataset called [CIFAR-10](https://www.cs.toronto.edu/~kriz/cifar.html).\n",
                "\n",
                "First, let's load CIFAR-10 into a Ray Dataset."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "d3f2e890",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to data/cifar-10-python.tar.gz\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "100%|██████████| 170498071/170498071 [00:21<00:00, 7792736.24it/s]\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting data/cifar-10-python.tar.gz to data\n",
                        "Files already downloaded and verified\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "2022-10-23 10:33:48,403\tINFO worker.py:1518 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
                    ]
                }
            ],
            "source": [
                "import ray\n",
                "import torchvision\n",
                "import torchvision.transforms as transforms\n",
                "\n",
                "train_dataset = torchvision.datasets.CIFAR10(\"data\", download=True, train=True)\n",
                "test_dataset = torchvision.datasets.CIFAR10(\"data\", download=True, train=False)\n",
                "\n",
                "train_dataset: ray.data.Dataset = ray.data.from_torch(train_dataset)\n",
                "test_dataset: ray.data.Dataset = ray.data.from_torch(test_dataset)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "a2e7db56",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "5d97a30cd75b40208a984ffa63cfecff",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "VBox(children=(HTML(value='<h2>Dataset</h2>'), Tab(children=(HTML(value='<div class=\"scrollableTable jp-Render…"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "train_dataset"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a89b59e8",
            "metadata": {},
            "source": [
                "{py:class}`from_torch <ray.data.from_torch>` doesn't parallelize reads, so you shouldn't use it with larger datasets.\n",
                "\n",
                "Next, let's represent our data using a dictionary of ndarrays instead of tuples. This lets us call {py:meth}`Dataset.iter_torch_batches <ray.data.Dataset.iter_torch_batches>` later in the tutorial."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "id": "9c485ff8",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Read->Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(_map_block_nosplit pid=3958)\u001b[0m Files already downloaded and verified\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Read->Map_Batches: 100%|██████████| 1/1 [00:04<00:00,  4.27s/it]\n",
                        "Read->Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(_map_block_nosplit pid=3958)\u001b[0m Files already downloaded and verified\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Read->Map_Batches: 100%|██████████| 1/1 [00:01<00:00,  1.40s/it]\n"
                    ]
                }
            ],
            "source": [
                "from typing import Dict, Tuple\n",
                "import numpy as np\n",
                "from PIL.Image import Image\n",
                "import torch\n",
                "\n",
                "\n",
                "def convert_batch_to_numpy(batch: Tuple[Image, int]) -> Dict[str, np.ndarray]:\n",
                "    images = np.stack([np.array(image) for image, _ in batch])\n",
                "    labels = np.array([label for _, label in batch])\n",
                "    return {\"image\": images, \"label\": labels}\n",
                "\n",
                "\n",
                "train_dataset = train_dataset.map_batches(convert_batch_to_numpy)\n",
                "test_dataset = test_dataset.map_batches(convert_batch_to_numpy)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "id": "4b416bbb",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "45e775b230614686909186cd0a845e37",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "VBox(children=(HTML(value='<h2>Dataset</h2>'), Tab(children=(HTML(value='<div class=\"scrollableTable jp-Render…"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "train_dataset"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a2cc4086",
            "metadata": {},
            "source": [
                "## Train a convolutional neural network\n",
                "\n",
                "Now that we've created our datasets, let's define the training logic."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "id": "ece6f85f",
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "import torch.nn as nn\n",
                "import torch.nn.functional as F\n",
                "\n",
                "\n",
                "class Net(nn.Module):\n",
                "    def __init__(self):\n",
                "        super().__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, 120)\n",
                "        self.fc2 = nn.Linear(120, 84)\n",
                "        self.fc3 = nn.Linear(84, 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 = torch.flatten(x, 1)  # flatten all dimensions except batch\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": "f9d85c9d",
            "metadata": {},
            "source": [
                "We define our training logic in a function called `train_loop_per_worker`. This function contains regular PyTorch code with a few notable exceptions:\n",
                "* We wrap our model with {py:func}`train.torch.prepare_model <ray.train.torch.prepare_model>`.\n",
                "* We call {py:func}`session.get_dataset_shard <ray.air.session.get_dataset_shard>` and {py:meth}`Dataset.iter_torch_batches <ray.data.Dataset.iter_torch_batches>` to get a subset of our training data.\n",
                "* We save model state using {py:func}`session.report <ray.air.session.report>`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "id": "6d32d183",
            "metadata": {},
            "outputs": [],
            "source": [
                "from ray import train\n",
                "from ray.air import session, Checkpoint\n",
                "from ray.train.torch import TorchCheckpoint\n",
                "import torch.nn as nn\n",
                "import torch.optim as optim\n",
                "import torchvision\n",
                "\n",
                "\n",
                "def train_loop_per_worker(config):\n",
                "    model = train.torch.prepare_model(Net())\n",
                "\n",
                "    criterion = nn.CrossEntropyLoss()\n",
                "    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
                "\n",
                "    train_dataset_shard = session.get_dataset_shard(\"train\")\n",
                "\n",
                "    for epoch in range(2):\n",
                "        running_loss = 0.0\n",
                "        train_dataset_batches = train_dataset_shard.iter_torch_batches(\n",
                "            batch_size=config[\"batch_size\"],\n",
                "        )\n",
                "        for i, batch in enumerate(train_dataset_batches):\n",
                "            # get the inputs and labels\n",
                "            inputs, labels = batch[\"image\"], batch[\"label\"]\n",
                "\n",
                "            # zero the parameter gradients\n",
                "            optimizer.zero_grad()\n",
                "\n",
                "            # forward + backward + optimize\n",
                "            outputs = model(inputs)\n",
                "            loss = criterion(outputs, labels)\n",
                "            loss.backward()\n",
                "            optimizer.step()\n",
                "\n",
                "            # print statistics\n",
                "            running_loss += loss.item()\n",
                "            if i % 2000 == 1999:  # print every 2000 mini-batches\n",
                "                print(f\"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}\")\n",
                "                running_loss = 0.0\n",
                "\n",
                "        metrics = dict(running_loss=running_loss)\n",
                "        checkpoint = TorchCheckpoint.from_state_dict(model.state_dict())\n",
                "        session.report(metrics, checkpoint=checkpoint)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "76f83b27",
            "metadata": {},
            "source": [
                "To improve our model's accuracy, we'll also define a `Preprocessor` to normalize the images."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f25ced31",
            "metadata": {},
            "outputs": [],
            "source": [
                "from ray.data.preprocessors import TorchVisionPreprocessor\n",
                "\n",
                "transform = transforms.Compose(\n",
                "    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
                ")\n",
                "preprocessor = TorchVisionPreprocessor(columns=[\"image\"], transform=transform)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "58100f87",
            "metadata": {},
            "source": [
                "Finally, we can train our model. This should take a few minutes to run."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "id": "89a51244",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/html": [
                            "== Status ==<br>Current time: 2022-08-30 15:31:37 (running for 00:00:45.17)<br>Memory usage on this node: 16.9/32.0 GiB<br>Using FIFO scheduling algorithm.<br>Resources requested: 0/10 CPUs, 0/0 GPUs, 0.0/14.83 GiB heap, 0.0/2.0 GiB objects<br>Result logdir: /Users/bveeramani/ray_results/TorchTrainer_2022-08-30_15-30-52<br>Number of trials: 1/1 (1 TERMINATED)<br><table>\n",
                            "<thead>\n",
                            "<tr><th>Trial name              </th><th>status    </th><th>loc           </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  running_loss</th><th style=\"text-align: right;\">  _timestamp</th><th style=\"text-align: right;\">  _time_this_iter_s</th></tr>\n",
                            "</thead>\n",
                            "<tbody>\n",
                            "<tr><td>TorchTrainer_6799a_00000</td><td>TERMINATED</td><td>127.0.0.1:3978</td><td style=\"text-align: right;\">     2</td><td style=\"text-align: right;\">         43.7121</td><td style=\"text-align: right;\">       595.445</td><td style=\"text-align: right;\">  1661898697</td><td style=\"text-align: right;\">            20.8503</td></tr>\n",
                            "</tbody>\n",
                            "</table><br><br>"
                        ],
                        "text/plain": [
                            "<IPython.core.display.HTML object>"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m 2022-08-30 15:30:54,566\tINFO config.py:71 -- Setting up process group for: env:// [rank=0, world_size=2]\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m 2022-08-30 15:30:55,727\tINFO train_loop_utils.py:300 -- Moving model to device: cpu\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m 2022-08-30 15:30:55,728\tINFO train_loop_utils.py:347 -- Wrapping provided model in DDP.\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1,  2000] loss: 2.276\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1,  2000] loss: 2.270\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1,  4000] loss: 1.964\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1,  4000] loss: 1.936\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1,  6000] loss: 1.753\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1,  6000] loss: 1.754\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1,  8000] loss: 1.638\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1,  8000] loss: 1.661\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1, 10000] loss: 1.586\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1, 10000] loss: 1.547\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [1, 12000] loss: 1.489\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [1, 12000] loss: 1.476\n",
                        "Result for TorchTrainer_6799a_00000:\n",
                        "  _time_this_iter_s: 20.542800188064575\n",
                        "  _timestamp: 1661898676\n",
                        "  _training_iteration: 1\n",
                        "  date: 2022-08-30_15-31-16\n",
                        "  done: false\n",
                        "  experiment_id: c25700542bc348dbbeaf54e46f1fc84c\n",
                        "  hostname: MBP.local.meter\n",
                        "  iterations_since_restore: 1\n",
                        "  node_ip: 127.0.0.1\n",
                        "  pid: 3978\n",
                        "  running_loss: 687.5853321105242\n",
                        "  should_checkpoint: true\n",
                        "  time_since_restore: 22.880314111709595\n",
                        "  time_this_iter_s: 22.880314111709595\n",
                        "  time_total_s: 22.880314111709595\n",
                        "  timestamp: 1661898676\n",
                        "  timesteps_since_restore: 0\n",
                        "  training_iteration: 1\n",
                        "  trial_id: 6799a_00000\n",
                        "  warmup_time: 0.0025300979614257812\n",
                        "  \n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2,  2000] loss: 1.417\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2,  2000] loss: 1.431\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2,  4000] loss: 1.403\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2,  4000] loss: 1.404\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2,  6000] loss: 1.394\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2,  6000] loss: 1.368\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2,  8000] loss: 1.343\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2,  8000] loss: 1.363\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2, 10000] loss: 1.340\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2, 10000] loss: 1.297\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3980)\u001b[0m [2, 12000] loss: 1.253\n",
                        "\u001b[2m\u001b[36m(RayTrainWorker pid=3979)\u001b[0m [2, 12000] loss: 1.276\n",
                        "Result for TorchTrainer_6799a_00000:\n",
                        "  _time_this_iter_s: 20.850306034088135\n",
                        "  _timestamp: 1661898697\n",
                        "  _training_iteration: 2\n",
                        "  date: 2022-08-30_15-31-37\n",
                        "  done: false\n",
                        "  experiment_id: c25700542bc348dbbeaf54e46f1fc84c\n",
                        "  hostname: MBP.local.meter\n",
                        "  iterations_since_restore: 2\n",
                        "  node_ip: 127.0.0.1\n",
                        "  pid: 3978\n",
                        "  running_loss: 595.4451928250492\n",
                        "  should_checkpoint: true\n",
                        "  time_since_restore: 43.71214985847473\n",
                        "  time_this_iter_s: 20.831835746765137\n",
                        "  time_total_s: 43.71214985847473\n",
                        "  timestamp: 1661898697\n",
                        "  timesteps_since_restore: 0\n",
                        "  training_iteration: 2\n",
                        "  trial_id: 6799a_00000\n",
                        "  warmup_time: 0.0025300979614257812\n",
                        "  \n",
                        "Result for TorchTrainer_6799a_00000:\n",
                        "  _time_this_iter_s: 20.850306034088135\n",
                        "  _timestamp: 1661898697\n",
                        "  _training_iteration: 2\n",
                        "  date: 2022-08-30_15-31-37\n",
                        "  done: true\n",
                        "  experiment_id: c25700542bc348dbbeaf54e46f1fc84c\n",
                        "  experiment_tag: '0'\n",
                        "  hostname: MBP.local.meter\n",
                        "  iterations_since_restore: 2\n",
                        "  node_ip: 127.0.0.1\n",
                        "  pid: 3978\n",
                        "  running_loss: 595.4451928250492\n",
                        "  should_checkpoint: true\n",
                        "  time_since_restore: 43.71214985847473\n",
                        "  time_this_iter_s: 20.831835746765137\n",
                        "  time_total_s: 43.71214985847473\n",
                        "  timestamp: 1661898697\n",
                        "  timesteps_since_restore: 0\n",
                        "  training_iteration: 2\n",
                        "  trial_id: 6799a_00000\n",
                        "  warmup_time: 0.0025300979614257812\n",
                        "  \n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "2022-08-30 15:31:37,386\tINFO tune.py:758 -- Total run time: 45.32 seconds (45.16 seconds for the tuning loop).\n"
                    ]
                }
            ],
            "source": [
                "from ray.train.torch import TorchTrainer\n",
                "from ray.air.config import ScalingConfig\n",
                "\n",
                "trainer = TorchTrainer(\n",
                "    train_loop_per_worker=train_loop_per_worker,\n",
                "    train_loop_config={\"batch_size\": 2},\n",
                "    datasets={\"train\": train_dataset},\n",
                "    scaling_config=ScalingConfig(num_workers=2),\n",
                "    preprocessor=preprocessor\n",
                ")\n",
                "result = trainer.fit()\n",
                "latest_checkpoint = result.checkpoint"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "1df4faa9",
            "metadata": {},
            "source": [
                "To scale your training script, create a [Ray Cluster](cluster-index) and increase the number of workers. If your cluster contains GPUs, add `\"use_gpu\": True` to your scaling config.\n",
                "\n",
                "```{code-block} python\n",
                "scaling_config=ScalingConfig(num_workers=8, use_gpu=True)\n",
                "```\n",
                "\n",
                "## Test the network on the test data\n",
                "\n",
                "Let's see how our model performs.\n",
                "\n",
                "To classify images in the test dataset, we'll need to create a {py:class}`Predictor <ray.train.predictor.Predictor>`.\n",
                "\n",
                "{py:class}`Predictors <ray.train.predictor.Predictor>` load data from checkpoints and efficiently perform inference. In contrast to {py:class}`TorchPredictor <ray.train.torch.TorchPredictor>`, which performs inference on a single batch, {py:class}`BatchPredictor <ray.train.batch_predictor.BatchPredictor>` performs inference on an entire dataset. Because we want to classify all of the images in the test dataset, we'll use a {py:class}`BatchPredictor <ray.train.batch_predictor.BatchPredictor>`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "id": "990ec534",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:01<00:00,  1.59s/it]\n"
                    ]
                }
            ],
            "source": [
                "from ray.train.torch import TorchPredictor\n",
                "from ray.train.batch_predictor import BatchPredictor\n",
                "\n",
                "batch_predictor = BatchPredictor.from_checkpoint(\n",
                "    checkpoint=latest_checkpoint,\n",
                "    predictor_cls=TorchPredictor,\n",
                "    model=Net(),\n",
                ")\n",
                "\n",
                "outputs: ray.data.Dataset = batch_predictor.predict(\n",
                "    data=test_dataset,\n",
                "    dtype=torch.float,\n",
                "    feature_columns=[\"image\"],\n",
                "    keep_columns=[\"label\"],\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "d20fd044",
            "metadata": {},
            "source": [
                "Our model outputs a list of energies for each class. To classify an image, we\n",
                "choose the class that has the highest energy."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "id": "00c8a336",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 59.42it/s]"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "{'prediction': 3, 'label': 3}\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\n"
                    ]
                }
            ],
            "source": [
                "import numpy as np\n",
                "\n",
                "\n",
                "def convert_logits_to_classes(df):\n",
                "    best_class = df[\"predictions\"].map(lambda x: x.argmax())\n",
                "    df[\"prediction\"] = best_class\n",
                "    return df[[\"prediction\", \"label\"]]\n",
                "\n",
                "\n",
                "predictions = outputs.map_batches(convert_logits_to_classes)\n",
                "\n",
                "predictions.show(1)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8973efc6",
            "metadata": {},
            "source": [
                "Now that we've classified all of the images, let's figure out which images were\n",
                "classified correctly. The ``predictions`` dataset contains predicted labels and \n",
                "the ``test_dataset`` contains the true labels. To determine whether an image \n",
                "was classified correctly, we join the two datasets and check if the predicted \n",
                "labels are the same as the actual labels."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "id": "8e6233ba",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 132.06it/s]"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "{'prediction': 3, 'label': 3, 'correct': True}\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\n"
                    ]
                }
            ],
            "source": [
                "def calculate_prediction_scores(df):\n",
                "    df[\"correct\"] = df[\"prediction\"] == df[\"label\"]\n",
                "    return df\n",
                "\n",
                "\n",
                "scores = predictions.map_batches(calculate_prediction_scores)\n",
                "\n",
                "scores.show(1)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8d401d91",
            "metadata": {},
            "source": [
                "To compute our test accuracy, we'll count how many images the model classified \n",
                "correctly and divide that number by the total number of test images."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "id": "29b2e2c2",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Shuffle Map: 100%|██████████| 1/1 [00:00<00:00, 152.00it/s]\n",
                        "Shuffle Reduce: 100%|██████████| 1/1 [00:00<00:00, 219.54it/s]\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "0.557"
                        ]
                    },
                    "execution_count": 13,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "scores.sum(on=\"correct\") / scores.count()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f0c84e2c",
            "metadata": {},
            "source": [
                "## Deploy the network and make a prediction\n",
                "\n",
                "Our model seems to perform decently, so let's deploy the model to an \n",
                "endpoint. This allows us to make predictions over the Internet."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "id": "f2faaf4c",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(ServeController pid=3987)\u001b[0m INFO 2022-08-30 15:31:39,948 controller 3987 http_state.py:129 - Starting HTTP proxy with name 'SERVE_CONTROLLER_ACTOR:SERVE_PROXY_ACTOR-4b114e48c80d3549aa5da89fa16707e0334a0bafde984fd8b8618e47' on node '4b114e48c80d3549aa5da89fa16707e0334a0bafde984fd8b8618e47' listening on '127.0.0.1:8000'\n",
                        "\u001b[2m\u001b[36m(HTTPProxyActor pid=3988)\u001b[0m INFO:     Started server process [3988]\n",
                        "\u001b[2m\u001b[36m(ServeController pid=3987)\u001b[0m INFO 2022-08-30 15:31:40,567 controller 3987 deployment_state.py:1232 - Adding 1 replica to deployment 'PredictorDeployment'.\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "RayServeSyncHandle(deployment='PredictorDeployment')"
                        ]
                    },
                    "execution_count": 14,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "from ray import serve\n",
                "from ray.serve import PredictorDeployment\n",
                "from ray.serve.http_adapters import json_to_ndarray\n",
                "\n",
                "\n",
                "serve.run(\n",
                "    PredictorDeployment.bind(\n",
                "        TorchPredictor,\n",
                "        latest_checkpoint,\n",
                "        model=Net(),\n",
                "        http_adapter=json_to_ndarray,\n",
                "    )\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "90327e8a",
            "metadata": {},
            "source": [
                "Let's classify a test image."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "id": "40da3863",
            "metadata": {},
            "outputs": [],
            "source": [
                "image = test_dataset.take(1)[0][\"image\"]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "556d94b7",
            "metadata": {},
            "source": [
                "You can perform inference against a deployed model by posting a dictionary with an `\"array\"` key. To learn more about the default input schema, read the {py:class}`NdArray <ray.serve.http_adapters.NdArray>` documentation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "id": "45bd85d6",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "[-1.1342155933380127,\n",
                            " -1.854529857635498,\n",
                            " 1.2062205076217651,\n",
                            " 2.6219608783721924,\n",
                            " 0.5199968218803406,\n",
                            " 2.2016565799713135,\n",
                            " 0.9447429180145264,\n",
                            " -0.5387609004974365,\n",
                            " -1.9515650272369385,\n",
                            " -1.676588773727417]"
                        ]
                    },
                    "execution_count": 16,
                    "metadata": {},
                    "output_type": "execute_result"
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\u001b[2m\u001b[36m(HTTPProxyActor pid=3988)\u001b[0m INFO 2022-08-30 15:31:41,713 http_proxy 127.0.0.1 http_proxy.py:315 - POST / 200 12.9ms\n",
                        "\u001b[2m\u001b[36m(ServeReplica:PredictorDeployment pid=3995)\u001b[0m INFO 2022-08-30 15:31:41,712 PredictorDeployment PredictorDeployment#pTPSPE replica.py:482 - HANDLE __call__ OK 9.9ms\n"
                    ]
                }
            ],
            "source": [
                "import requests\n",
                "\n",
                "payload = {\"array\": image.tolist(), \"dtype\": \"float32\"}\n",
                "response = requests.post(\"http://localhost:8000/\", json=payload)\n",
                "response.json()"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3.8.10 ('venv': venv)",
            "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.8.10"
        },
        "vscode": {
            "interpreter": {
                "hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}