04_Scaling_model_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Scaling model training\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 guides you through distributed model training with Ray. Through fine-tuning a transformer for a computer vision task, ML practitioners will learn how to scale training workloads using deep learning models on large datasets.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "For this notebook, you should satisfy the following minimum requirements:\n",
    "-   Practical Python knowledge.\n",
    "-   Familiarity with training deep learning models.\n",
    "-   Experience with Ray equivalent to completing the following 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",
    "    \n",
    "### Learning objectives\n",
    "\n",
    "-   Understand the challenges associated with distributing model training across multiple GPUs.\n",
    "-   Implement the data parallelism design pattern using Ray Datasets.\n",
    "-   Fine-tune a transformer model on an image dataset using Ray Train.\n",
    "\n",
    "### What will you do?\n",
    "\n",
    "-   Distributed model training overview\n",
    "    -   Learn about why training large machine learning models requires a distributed solution.\n",
    "    -   Refresh your knowledge of the data parallelism design pattern.\n",
    "-   Example: Fine-tuning a model for image segmentation.\n",
    "    -   Background\n",
    "        -   Data - MITADE20K benchmark dataset of scene images.\n",
    "        -   Model - Segformer transformer for semantic segmentation.\n",
    "    -   Getting started\n",
    "        -   Start Ray cluster and set-up environment.\n",
    "    -   Data ingest\n",
    "        -   Batch and transform raw data into training inputs using Ray Data.\n",
    "    -   Distributed training\n",
    "        -   Fine-tune transformer model on benchmark dataset using Ray Train.\n",
    "-   Conclusion\n",
    "    -   Summarize the distributed training approach as well as the Ray components at each stage of the pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed model training\n",
    "\n",
    "As the development of machine learning models advances, their [size continues to balloon](https://epochai.org/blog/machine-learning-model-sizes-and-the-parameter-gap). Training these large neural networks can take a prohibitively long time and requires an increasingly [massive amount of compute](https://www.hyro.ai/glossary/gpt-3#:~:text=To%20be%20exact%2C%20GPT%2D3,amount%20of%20time%20is%20unimaginable.).\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Overview_of_Ray/ai_compute_annotated.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|OpenAI's blog[\"AI and compute\"](https://openai.com/research/ai-and-compute) reports that the amount of compute needed to train the largest models has roughly doubled every 3.5 months since 2012, with no signs of this trend slowing down. Annotated original chart with trend lines overlaid.|\n",
    "\n",
    "Distributing this workload presents unique challenges involved with orchestrating multiple machines to produce one computationally synchronized result. This problem only compounds when working with heterogeneous resources, multiple tuning experiments, or a model that can't fit on a single GPU. To address these issues, machine learning practitioners have developed a variety of techniques to parallelize training across nodes, one of which is data parallelism.\n",
    "\n",
    "### Data parallelism\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Data parallelism:</a></strong> a design pattern that trains replicas of the model on different subsets of a large dataset, periodically synchronizing weights to produce a fully trained result. This method requires that a model's parameters, or weights, are able to fit on a single GPU's memory.\n",
    "</div>\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Scaling_model_training/data_parallelism.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|A large dataset is sharded across multiple worker nodes each containing a model copy. Gradients calculated on independent nodes are continuously synchronized with others to produce a final trained model.|\n",
    "\n",
    "Ray Train provides distributed data parallel training capabilities. Its integration with Ray AIR also allows for convenient parallelization of data ingestion and pre-processing, hyperparameter tuning, batch inference, and serving. This provides a unified compute layer for the machine learning pipeline, eliminating the need to stitch together independent scaling solutions at each stage. In the next section, you will implement this design pattern using a transformer model and scene images to accomplish a computer vision task.\n",
    "\n",
    "Note: There are other techniques for distributed training such as model parallelism which divides the model itself across multiple GPUs. However, this module will focus on implementing data parallelism."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background on semantic segmentation\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Semantic segmentation:</a></strong> a computer vision task that assigns labels to object regions in a scene, pixel-by-pixel. Similar to object detection, this approach involves dividing an image into multiple semantic categories such as couch, person, car, or sky.\n",
    "</div>\n",
    "\n",
    "In this hands-on example, you will implement the data parallelism design pattern by fine-tuning a pretrained transformer model on scene image data.\n",
    "\n",
    "### Data\n",
    "\n",
    "#### MIT ADE20K - scene parsing benchmark\n",
    "\n",
    "The [MIT ADE20K Dataset](http://sceneparsing.csail.mit.edu/) (also known as \"SceneParse150\") provides the largest open source dataset for scene parsing. It is often used as a standard for assessing semantic segmentation model performance due to its high-quality annotations.\n",
    "\n",
    "You will use the training set for fine-tuning and the unlabeled test set for evaluation.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Scaling_inference/scene.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Test image on the left vs. predicted result. ([Source](https://github.com/CSAILVision/semantic-segmentation-pytorch))|\n",
    "\n",
    "Dataset highlights\n",
    "\n",
    "-   20k annotated, scene-centric training images\n",
    "-   3.3k unlabeled test images\n",
    "-   150 [semantic categories](https://docs.google.com/spreadsheets/d/1se8YEtb2detS7OuPE86fXGyD269pMycAWe2mtKUj2W8/edit?usp=sharing) (such as person, car, bed, sky, etc.)\n",
    "\n",
    "### Model\n",
    "\n",
    "#### SegFormer - transformer-based framework for semantic segmentation\n",
    "\n",
    "[SegFormer](https://arxiv.org/pdf/2105.15203.pdf) is an effective semantic segmentation method based on a transformer architecture. [Transformers](https://en.wikipedia.org/wiki/Transformer_(machine_learning_model)) are a type of deep learning architecture that process sequential data via a series of self-attention layers and then transform them via a feedforward neural network.\n",
    "\n",
    "What sets SegFormer apart from previous transformer-based approaches are two key features:\n",
    "\n",
    "1.  A hierarchically structured transformer encoder which does not depend on positional encoding that avoids interpolation when training and testing resolutions differ.\n",
    "2.  A lightweight MLP layer that avoids complex decoders.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Scaling_inference/segformer_architecture.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Segformer architecture illustrated in the [original paper](https://arxiv.org/pdf/2105.15203.pdf).|\n",
    "\n",
    "You will use a general, pre-trained SegFormer model to fine-tune on [MITADE20K](http://sceneparsing.csail.mit.edu/) image data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up necessary imports and utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "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)\n",
    "\n",
    "warnings.simplefilter(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Ray runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ray.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the model components from the HuggingFace Hub\n",
    "\n",
    "From the [Hugging Face Hub](https://huggingface.co/docs/hub/index), retrieve the pretrained SegFormer model by specifying the model name and [label files](https://huggingface.co/datasets/huggingface/label-files/blob/main/ade20k-id2label.json) which map indices to semantic categories."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load label mappings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import get_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "id2label, label2id = get_labels()\n",
    "num_labels = len(id2label)\n",
    "\n",
    "print(f\"Total number of labels: {len(id2label)}\")\n",
    "print(f\"Example labels: {list(id2label.values())[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The utility function `get_labels` fetches two dictionary mappings from [Hugging Face](https://huggingface.co/datasets/huggingface/label-files/blob/main/ade20k-id2label.json), `id2label` and `label2id`, which are used to convert between numerical and string labels for the 150 available [semantic categories](https://docs.google.com/spreadsheets/d/1se8YEtb2detS7OuPE86fXGyD269pMycAWe2mtKUj2W8/edit#gid=0) of objects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load SegFormer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import SegformerForSemanticSegmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"nvidia/mit-b0\" # https://huggingface.co/nvidia/mit-b0\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": [
    "The [Hugging Face Hub](https://huggingface.co/nvidia/mit-b0) makes available many variations on SegFormer. Here, you specify a version pretrained on the Imagenet-1K dataset. We will be further fine-tuning it on the MITADE20K (SceneParse150) dataset in this notebook.\n",
    "\n",
    "Note: This \"b0\" model is the smallest, with [other options](https://huggingface.co/nvidia/mit-b0) ranging up to and including \"b5\". Keep this in mind as something to experiment with when comparing different batch inference architectures later on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data ingest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset from HuggingFace Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from utils import convert_image_to_rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SMALL_DATA = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "  <strong>SMALL_DATA</strong>: a flag to download a subset (160 images) of the available data. Defaults to True. Set to False (recommended) to work with the full train data (20k images).\n",
    "</div>\n",
    "\n",
    "If you set `SMALL_DATA` to `False`, expect it to take some time (depending on your connection download speed) because you are downloading all test images to your local machine or cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load SceneParse150"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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[:160]\")\n",
    "else:\n",
    "    train_dataset = load_dataset(DATASET_NAME, split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.map(convert_image_to_rgb)\n",
    "train_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each sample contains three components:\n",
    "* **`image`** \n",
    "    * The PIL image.\n",
    "* **`annotation`**  \n",
    "    * Human annotations of image regions (annotation mask is `None` in testing set).\n",
    "* **`category`**  \n",
    "    * Category of the scene generally (e.g. driveway, voting booth, dairy_outdoor)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Display example images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import display_example_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Try running this multiple times!\n",
    "display_example_images(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Ray Dataset for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 8\n",
    "N_BATCHES = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import get_image_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get BATCH_SIZE * N_BATCHES randomly shuffled image IDs from the train dataset.\n",
    "image_indices = get_image_indices(dataset=train_dataset, n=BATCH_SIZE * N_BATCHES)\n",
    "\n",
    "# Create a list of tuples (image, label) for the indices sampled from the train dataset.\n",
    "data = [\n",
    "    (train_dataset[i][\"image\"], train_dataset[i][\"annotation\"]) for i in image_indices\n",
    "]\n",
    "\n",
    "# Create a Ray Dataset from the list of images to use in Ray AIR.\n",
    "train_ds = ray.data.from_items(data)\n",
    "train_ds = train_ds.map_batches(\n",
    "    lambda x: pd.DataFrame(x, columns=[\"image\", \"annotation\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_ds.schema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Display example image\n",
    "train_ds.take(1)[0][\"image\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Display example image\n",
    "train_ds.take(1)[0][\"annotation\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create preprocessor for distributed data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import SegformerImageProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def images_preprocessor(batch):\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    segformer_image_processor = SegformerImageProcessor.from_pretrained(\n",
    "        MODEL_NAME, do_reduce_labels=True\n",
    "    )\n",
    "\n",
    "    # inputs are `transformers.image_processing_utils.BatchFeature`\n",
    "    inputs = segformer_image_processor(\n",
    "        images=list(batch[\"image\"]),\n",
    "        segmentation_maps=list(batch[\"annotation\"]),\n",
    "        return_tensors=\"np\",\n",
    "    )\n",
    "\n",
    "    return dict(inputs)  # {\"pixel_values\": array, \"labels\": array}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Feature extractors](https://huggingface.co/docs/transformers/main_classes/feature_extractor) preprocess input features (e.g. image data) by normalizing, resizing, padding, and converting raw images into the shape expected by SegFormer.\n",
    "\n",
    "The [`reduce_labels`](https://huggingface.co/docs/transformers/model_doc/segformer#segformer) flag ensures that the background of an image (anything that is not explicitly an object) isn't included when computing loss. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ray.data.preprocessors import BatchMapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_preprocessor = BatchMapper(\n",
    "    fn=images_preprocessor, batch_format=\"pandas\", batch_size=2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run distributed training with SegFormer from HuggingFace you need:\n",
    "\n",
    "* setup batches preprocessor\n",
    "* setup HuggingFace Trainer configuration for all workers\n",
    "* create HuggingFaceTrainer - Ray Train object that handles distributed training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup HuggingFace Trainer per worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def trainer_init_per_worker(train_dataset, eval_dataset=None, **config):\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "\n",
    "    name = \"segformer-finetuned\"\n",
    "\n",
    "    # Setup model\n",
    "    segformer = SegformerForSemanticSegmentation.from_pretrained(\n",
    "        MODEL_NAME, id2label=id2label, label2id=label2id\n",
    "    )\n",
    "\n",
    "    # Setup optimizer and LR scheduler\n",
    "    optimizer = torch.optim.AdamW(params=segformer.parameters(), lr=1e-4)\n",
    "    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(\n",
    "        optimizer=optimizer, lr_lambda=lambda x: x\n",
    "    )\n",
    "\n",
    "    # Setup HF Training Arguments\n",
    "    training_args = TrainingArguments(\n",
    "        name,\n",
    "        num_train_epochs=5,\n",
    "        per_device_train_batch_size=BATCH_SIZE,\n",
    "        save_total_limit=3,\n",
    "        save_strategy=\"epoch\",\n",
    "        logging_strategy=\"epoch\",\n",
    "        eval_accumulation_steps=2,\n",
    "        log_level=\"error\",\n",
    "        log_level_replica=\"error\",\n",
    "        log_on_each_node=False,\n",
    "        remove_unused_columns=False,\n",
    "        push_to_hub=False,\n",
    "        disable_tqdm=True,  # declutter the output a little\n",
    "        no_cuda=True,\n",
    "    )\n",
    "\n",
    "    # Setup HF Trainer\n",
    "    trainer = Trainer(\n",
    "        model=segformer,\n",
    "        optimizers=(optimizer, lr_scheduler),\n",
    "        args=training_args,\n",
    "        train_dataset=train_dataset,\n",
    "    )\n",
    "\n",
    "    print(\"Starting training...\")\n",
    "    return trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create HuggingFace Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ray.train.huggingface import HuggingFaceTrainer\n",
    "from ray.air.config import RunConfig, ScalingConfig, CheckpointConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# setup parameters for the ScalingConfig\n",
    "num_workers = 2\n",
    "use_gpu = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup Ray's HF Trainer\n",
    "trainer = HuggingFaceTrainer(\n",
    "    trainer_init_per_worker=trainer_init_per_worker,\n",
    "    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
    "    datasets={\n",
    "        \"train\": train_ds,\n",
    "    },\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        ),\n",
    "    ),\n",
    "    preprocessor=batch_preprocessor,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run model training\n",
    "result = trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercises**\n",
    "\n",
    "Try the following modifications to our work here\n",
    "- Replace the existing model with another SegFormer, described here: https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512\n",
    "- Switch over to using GPUs for training\n",
    "- If you have time and want to see how a raw PyTorch model is integrated with Ray Train, try the example at https://docs.ray.io/en/latest/train/examples/pytorch/torch_fashion_mnist_example.html\n",
    "    - hint: you may need to update `torchvision` to get this example to run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Congratulations! You have successfully gained experience in using Ray Train to fine-tune a vision transformer model for semantic segmentation. In the upcoming module, you will be able to build on this example by performing distributed batch inference to inspect the results of the model.\n",
    "\n",
    "### Summary\n",
    "\n",
    "-   Distributed model training\n",
    "    -   Training and fine-tuning large neural networks requires a massive amount of compute, so the only solution is to distribute this workload.\n",
    "    -   Data parallelism offers a pattern for sharding a large dataset across multiple machines for training and gradient synchronization.\n",
    "    -   This orchestration and maintenance is challenging, and Ray AIR offers a unified compute solution to scale this workload that integrates well with other stages in the pipeline.\n",
    "-   Fine-tuning Segformer on MITADE20K\n",
    "    -   Data ingest\n",
    "        -   Ray Data can be used to ingest and preprocess training images. These same transformations can be applied during tuning, inference, and serving.\n",
    "    -   Distributed training\n",
    "        -   Ray Train can fine-tune a transformer model, in this case implementing the data parallel design pattern by running PyTorch's [Distributed Data Parallel](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) as the backend.\n"
   ]
  }
 ],
 "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.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}