05_Batch_inference_with_Ray_Datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Batch inference with Ray Datasets\n",
    "\n",
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Generic/ray_logo.png\" width=\"20%\" loading=\"lazy\">\n",
    "\n",
    "## About this notebook\n",
    "\n",
    "### Is this module right for you?\n",
    "\n",
    "This module is an extension of [Scaling Batch Inference](https://github.com/ray-project/ray-educational-materials/blob/main/Computer_vision_workloads/Semantic_segmentation/Scaling_batch_inference.ipynb) and presents another approach for distributed batch inference on Ray. Through this short exploration of Ray Datasets' `map_batches` functionality, you will approach the same semantic segmentation task as before and use only Datasets to generate predictions.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "For this notebook you should satisfy the following requirements:\n",
    "\n",
    "* Practical Python and machine learning experience.\n",
    "* Familiarity with batch inference in ML.\n",
    "* Familiarity with Ray and Ray AIR equivalent to completing these training modules:\n",
    "  * [Overview of Ray](https://github.com/ray-project/ray-educational-materials/blob/main/Introductory_modules/Overview_of_Ray.ipynb)\n",
    "  * [Introduction to Ray AIR](https://github.com/ray-project/ray-educational-materials/blob/main/Introductory_modules/Introduction_to_Ray_AIR.ipynb)\n",
    "  * [Ray Core](https://github.com/ray-project/ray-educational-materials/tree/main/Ray_Core)\n",
    "\n",
    "### Learning objectives\n",
    "\n",
    "* Implement batch inference on a semantic segmentation task with Ray Datasets by using `map_batches`.\n",
    "* Customize the compute resources for inference by experimenting with the `ActorPoolStrategy`.\n",
    "\n",
    "### What will you do?\n",
    "\n",
    "* Set up environment.\n",
    "* Distributed batch inference with Ray Datasets.\n",
    "  * Create a Ray Dataset.\n",
    "  * Create a prediction class with inference logic.\n",
    "  * (Optional) Specify compute by defining an actor pool strategy.\n",
    "  * Use `map_batches` to apply the prediction class on batches to perform inference.\n",
    "* Experiment with compute resources and observability in the coding exercises."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up environment from \"Scaling Batch Inference\" module\n",
    "\n",
    "This notebook provides another approach to batch inference on semantic segmentation tasks. In order to extend the solution posed in the main module, you must re-establish the context of the original example. For more context regarding these steps, refer to the root [Scaling Batch Inference](https://github.com/ray-project/ray-educational-materials/blob/main/Computer_vision_workloads/Semantic_segmentation/Scaling_batch_inference.ipynb) notebook.\n",
    "\n",
    "In this section, you will import and load in the necessary task components:\n",
    "\n",
    "* Set up necessary imports and utilities.\n",
    "* Load label mappings.\n",
    "* Load SegFormer.\n",
    "* Load the feature extractor.\n",
    "* Load the dataset.\n",
    "\n",
    "In addition, you will port over some Ray-specific actions:\n",
    "\n",
    "* Initialize Ray runtime.\n",
    "* Put the model and feature extractor in the object store."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up necessary imports and utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "from PIL.JpegImagePlugin import JpegImageFile\n",
    "\n",
    "# Set the seed to a fixed value for reproducibility.\n",
    "torch.manual_seed(201)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load label mappings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id2label, label2id = get_labels()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load SegFormer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import SegformerForSemanticSegmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"nvidia/segformer-b0-finetuned-ade-512-512\"\n",
    "\n",
    "segformer = SegformerForSemanticSegmentation.from_pretrained(\n",
    "    MODEL_NAME, id2label=id2label, label2id=label2id\n",
    ")\n",
    "\n",
    "print(f\"Number of model parameters: {segformer.num_parameters()/(10**6):.2f} M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the feature extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import SegformerFeatureExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segformer_feature_extractor = SegformerFeatureExtractor.from_pretrained(\n",
    "    MODEL_NAME, reduce_labels=True\n",
    ")\n",
    "segformer_feature_extractor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from utils import convert_image_to_rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SMALL_DATA = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "  <strong>SMALL_DATA</strong>: default `True` - set to download only 160 images from the data set. Set to `False` (recommended) to work with full testing dataset (3352 images).\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_NAME = \"scene_parse_150\"\n",
    "\n",
    "# Load data from the Hugging Face datasets repository.\n",
    "if SMALL_DATA:\n",
    "    train_dataset = load_dataset(DATASET_NAME, split=\"train[:10]\")\n",
    "    test_dataset = load_dataset(DATASET_NAME, split=\"test[:160]\")\n",
    "else:\n",
    "    train_dataset = load_dataset(DATASET_NAME, split=\"train[:10]\")\n",
    "    test_dataset = load_dataset(DATASET_NAME, split=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = test_dataset.map(convert_image_to_rgb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Ray runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "\n",
    "ray.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Put the model and feature extractor in the object store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segformer_ref = ray.put(segformer)\n",
    "segformer_feature_extractor_ref = ray.put(segformer_feature_extractor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed batch inference with Ray Datasets\n",
    "\n",
    "[Ray Datasets](https://docs.ray.io/en/latest/data/dataset.html) are the standard way to load and exchange data in Ray libraries and applications. They are designed to provide distributed loading, preprocessing, and transformations such as [maps](https://docs.ray.io/en/latest/data/api/dataset.html#ray.data.Dataset.map_batches), [global and grouped aggregations](https://docs.ray.io/en/latest/data/api/grouped_dataset.html#ray.data.grouped_dataset.GroupedDataset), and [shuffling operations](https://docs.ray.io/en/latest/data/api/dataset.html#ray.data.Dataset.random_shuffle). In this notebook, you will be leveraging Ray Datasets' [`map_batches`](https://docs.ray.io/en/latest/data/api/dataset.html#ray.data.Dataset.map_batches) method as a means to perform batch inference. \n",
    "\n",
    "You will encounter the following steps:\n",
    "\n",
    "1. Create a Ray Dataset.\n",
    "2. Create a prediction class with inference logic.\n",
    "3. (Optional) Specify compute by defining an actor pool strategy.\n",
    "4. Use `map_batches` to apply the prediction class on batches to perform inference.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Scaling_inference/ray_datasets.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Ray Datasets parallelize data loading, preprocessing, and batching. To perform inference, Datasets are able to call `map_batches` to apply a function (the prediction logic) to batches in parallel.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Ray Dataset with 160 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_image_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "N_BATCHES = 10\n",
    "\n",
    "# Get BATCH_SIZE * N_BATCHES randomly shuffled image IDs from the test dataset.\n",
    "image_indices = get_image_indices(dataset=test_dataset, n=BATCH_SIZE * N_BATCHES)\n",
    "\n",
    "# Create a list of images for the indices sampled from the test dataset.\n",
    "data = [test_dataset[i][\"image\"] for i in image_indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a Ray Dataset from the list of images.\n",
    "dataset = ray.data.from_items(data)\n",
    "dataset.show(limit=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a prediction class with inference logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PredictionClass:\n",
    "    # The constructor method initializes the class to load/cache the model and feature extractor.\n",
    "    def __init__(\n",
    "        self,\n",
    "        model: SegformerForSemanticSegmentation,\n",
    "        feature_extractor: SegformerFeatureExtractor,\n",
    "    ):\n",
    "        self.model = model\n",
    "        self.feature_extractor = feature_extractor\n",
    "\n",
    "    def __call__(self, batch: list[JpegImageFile]) -> list[np.ndarray]:\n",
    "\n",
    "        # Set the device on which PyTorch will run.\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.model.to(device)  # Move the model to specified device.\n",
    "        self.model.eval()  # Set the model in evaluation mode on test data.\n",
    "\n",
    "        # The feature extractor processes raw images.\n",
    "        inputs = self.feature_extractor(images=batch, return_tensors=\"pt\")\n",
    "\n",
    "        # The model is applied to input images in the inference step.\n",
    "        with torch.no_grad():\n",
    "            outputs = self.model(pixel_values=inputs.pixel_values.to(device))\n",
    "\n",
    "        # Post-process the output for display.\n",
    "        image_sizes = [image.size[::-1] for image in batch]\n",
    "        segmentation_maps_postprocessed = (\n",
    "            self.feature_extractor.post_process_semantic_segmentation(\n",
    "                outputs=outputs, target_sizes=image_sizes\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # Return list of segmentation maps detached from the computation graph.\n",
    "        return [j.detach().cpu().numpy() for j in segmentation_maps_postprocessed]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify compute by defining an actor pool strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.data import ActorPoolStrategy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Ray Datasets, transformations can either be carried out by Ray Tasks or Actors. While the default compute strategy uses Ray Tasks, you can specify an `ActorPoolStrategy` which dynamically [autoscales](https://docs.ray.io/en/latest/data/transforming-datasets.html#compute-strategy) the number of actors between a  `min` and `max` size to carry out the transforms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run parallel batch inference on a Ray Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dataset = dataset.map_batches(\n",
    "    PredictionClass,\n",
    "    batch_size=1,\n",
    "    num_gpus=0,\n",
    "    num_cpus=1,\n",
    "    compute=ActorPoolStrategy(min_size=1, max_size=2),\n",
    "    fn_constructor_args=(segformer, segformer_feature_extractor),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the Dataset `map_batches()` [function](https://docs.ray.io/en/latest/data/api/dataset.html#ray.data.Dataset.map_batches) to apply the model to the Dataset in parallel. You can specify the batch size, any resources, as well as any autoscaling options for the actor pool.\n",
    "\n",
    "Note: don't forget to pass `fn_constructor_args` to construct `PredictionClass`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dataset.take(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running inference, you can inspect predictions to probe the resulting image array. Notice that the resulting predictions dataset is, itself, a Ray Dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercises**\n",
    "\n",
    "- In this approach using Ray Datasets, you used an `ActorPoolStrategy` to set an upper and lower bound on the autoscaling of the actor pool.\n",
    "    - A natural experiment is to try toggling the `min_size` and `max_size` of the actor pool in `map_batches` to see the effect on runtime performance.\n",
    "    - To extend this exercise even further, open up your Ray Dashboard (linked when you called `ray.init()`) and see the dynamic autoscaling of the actor pool live.\n",
    "- Another experiment is to integrate your own code using a simpler form of `map_batches` where the only argument is a Python function\n",
    "    - In the present case, the function signature will look like this: `(batch: list[JpegImageFile]) -> list[np.ndarray]`\n",
    "    - Start with a `no_op` function that just returns the batch, and test that with `map_batches`\n",
    "    - Change the function to return a grayscale version of each image by performing a `.mean()` across the color channels (axis 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary: Distributed batch inference with Ray Datasets\n",
    "\n",
    "#### Key API elements\n",
    "* **`Datasets`**\n",
    "    * These are used to parallelize data loading, preprocessing, and exchanging data in Ray AIR.\n",
    "\n",
    "* **`map_batches`**\n",
    "    * This is a function to apply a transformation and/or model class to all batches. Can be used as a way to perform batch inference using only Ray Datasets without introducing other components of Ray AIR."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "567405a8058597909526349386224fe35dd047505a91307e44ed44be00113429"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}