examples/Introducing_TorchEval.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ymmr9GYSaPe"
      },
      "source": [
        "# 🎉 Introducing TorchEval! 🎉"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M3p7hvw6dkwS"
      },
      "source": [
        "### Fork this notebook on [Google Colab](https://colab.research.google.com/github/pytorch/torcheval/blob/main/examples/Introducing_TorchEval.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F_TukTozUeih"
      },
      "source": [
        "First lets install TorchEval from pypi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "_20AovPdS5Lj"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install torcheval-nightly"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jmrSNmXISecu"
      },
      "source": [
        "TorchEval is a library built for users who want highly performant implementations of common metrics to evaluate machine learning models. It also provides an easy to use interface for building custom metrics with the same toolkit. Building your metrics with TorchEval makes running distributed training loops with torch.distributed a breeze."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7rk1b9YiSnWT"
      },
      "source": [
        "# Using Metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0z30-tXqSi8u"
      },
      "source": [
        "Let's setup a simple sequential model here with one layer and RELU activation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zqg6mjhNSixI"
      },
      "source": [
        "We'll also define a function that runs some random data through the randomly initialized network and return some random labels with the network output.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "PBQjQuIVSX2x"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "BATCH_SIZE = 8\n",
        "INPUT_SIZE = 10\n",
        "NUM_CLASSES = 6\n",
        "\n",
        "model = torch.nn.Sequential(torch.nn.Linear(INPUT_SIZE, NUM_CLASSES), torch.nn.ReLU())\n",
        "\n",
        "def get_outputs_and_targets():\n",
        "    input = torch.rand(size=(BATCH_SIZE, INPUT_SIZE))\n",
        "    target = torch.randint(size=(BATCH_SIZE,), high=NUM_CLASSES)\n",
        "    outputs = model(input)\n",
        "    return outputs, target"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sRNEiOYfSwzx"
      },
      "source": [
        "### Functional Implementations ([Docs](https://pytorch.org/torcheval/main/torcheval.metrics.functional.html))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bah2rjf_SzcH"
      },
      "source": [
        "Now let's see how accurate our random model was using a functional implementation of multiclass accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ei_QpUmJSz3O",
        "outputId": "7d82403b-0afd-42bd-bb8e-7d5273b5f5af"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(0.1250)\n"
          ]
        }
      ],
      "source": [
        "from torcheval.metrics.functional import multiclass_accuracy\n",
        "\n",
        "outputs, target = get_outputs_and_targets()\n",
        "\n",
        "print(\n",
        "    multiclass_accuracy(outputs, target)\n",
        ")  # just computes the metric value for this batch of data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1UhQSDlfTEsg"
      },
      "source": [
        "### Class Based Implementations ([Docs](https://pytorch.org/torcheval/main/torcheval.metrics.html))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ryHi3e8TH-X"
      },
      "source": [
        "Now lets assume we have a few batches of data we want to run through and aggregate the total results. This can be done easily with our interfaces inheriting from the `Metric` class "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5itVteZmS1xQ",
        "outputId": "2493cbbe-edfd-4248-a63b-6abc98a3e15d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(0.1875)\n"
          ]
        }
      ],
      "source": [
        "from torcheval.metrics import MulticlassAccuracy\n",
        "\n",
        "metric = MulticlassAccuracy()\n",
        "for i in range(10):\n",
        "    outputs, target = get_outputs_and_targets()\n",
        "    metric.update(outputs, target) # updates internal state variables with this batch of data \n",
        "                                   # those state variables will be used to compute the metric on all 10 batches of data\n",
        "\n",
        "print(metric.compute()) # compute metric on all 10 batches. Call metric.reset() to clear internal variables \n",
        "                        # and compute metric on a new set of data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "auOGBBV3dkwV"
      },
      "source": [
        "Deferred computation of the metric is a useful trick for speeding up evaluation loops. In some cases, accumulating state can be done quickly, while metric computation is slow.\n",
        "\n",
        "Now let's reset the internal state variables of our metric. Any following calls to `.compute()` will only compute the metric over data passed in since the last `.reset()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BRbVN7_ldkwV",
        "outputId": "91035011-f8d2-4350-b2ec-8e8788b85e24"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(1.)\n"
          ]
        }
      ],
      "source": [
        "metric.reset()\n",
        "metric.update(torch.tensor([1,0,1]), torch.tensor([1,0,1])) # should give 100% accuracy\n",
        "print(metric.compute())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VSLJOgneTLtc"
      },
      "source": [
        "### In A Distributed Setting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P86OH-rZTNfI"
      },
      "source": [
        "With distributed training loops, we we also need to synchronize results across processes. TorchEval makes this very simple, let's take a look at a basic training loop."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fvCMEu6dTOMQ"
      },
      "source": [
        "We have some boilerplate to get distributed training running"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "WMMk2-xYTJ5r"
      },
      "outputs": [],
      "source": [
        "import torch.distributed as dist\n",
        "import torch.distributed.launcher as pet\n",
        "import uuid\n",
        "\n",
        "lc = pet.LaunchConfig(\n",
        "        min_nodes=1,\n",
        "        max_nodes=1,\n",
        "        nproc_per_node=4,\n",
        "        run_id=str(uuid.uuid4()),\n",
        "        rdzv_backend=\"c10d\",\n",
        "        rdzv_endpoint=\"localhost:0\",\n",
        "        max_restarts=0,\n",
        "        monitor_interval=1,\n",
        "    )\n",
        "\n",
        "lc.start_method = \"fork\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-6Y7dZ8fTZZO"
      },
      "source": [
        "Now let's setup a distributed loop to do inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "Meq8ClTVTYxl"
      },
      "outputs": [],
      "source": [
        "from torcheval.metrics.toolkit import sync_and_compute # import sync_and_compute from our toolkit to sync data between processes\n",
        "from torcheval.metrics import MulticlassF1Score\n",
        "\n",
        "def distributed_loop():\n",
        "    dist.init_process_group(backend=\"gloo\")\n",
        "    metric = MulticlassF1Score(num_classes=NUM_CLASSES)\n",
        "    for i in range(10):\n",
        "        input = torch.rand(size=(BATCH_SIZE, INPUT_SIZE))\n",
        "        target = torch.randint(size=(BATCH_SIZE,), high=NUM_CLASSES)\n",
        "        outputs = model(input).detach()\n",
        "        metric.update(outputs, target)\n",
        "\n",
        "    return sync_and_compute(metric) # include recipient_rank=\"all\" for each process to return computed metric"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oBt1xYciTV5X",
        "outputId": "8052bf87-f8c0-4e47-cb93-b5af46a4342e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{0: tensor(0.1625), 1: None, 2: None, 3: None}\n"
          ]
        }
      ],
      "source": [
        "batch_values = pet.elastic_launch(lc, entrypoint=distributed_loop)()\n",
        "print(batch_values) # process with label \"rank 0\" is using all 40 batches of data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BEvJjAO3dkwW"
      },
      "source": [
        "You can see that on rank 0, we got back a metric value which was computed using all the data across processes. If we had used `sync_and_compute(metric, recipient_rank=\"all\")`, then each process would return `tensor(0.1750)`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sKtnbeORdkwW"
      },
      "source": [
        "A more full example of a distributed training setup is in our [examples directory](https://github.com/pytorch/torcheval/blob/main/examples/distributed_example.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UoR1eTMrTg90"
      },
      "source": [
        "# Adding your own metric"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WONJW_5JTjy7"
      },
      "source": [
        "To add your own metric, you simply need to inherit from `Metric` and implement 4 methods.\n",
        "\n",
        "1. `__init__(self,)`: Defines the state variables \n",
        "2. `update(self, *args)`: Determines how to update the state variables with new data\n",
        "3. `compute(self,)`: Computes the metric from the state variables\n",
        "4. `merge_state(self, metrics)`: Describes how to merge the internal states when the metric objects from all processes are collected on a single device.\n",
        "\n",
        "Below, we implement the Kolmogorov-Smirnov 2 Sample test, utilizing the implemenation already in `scipy`.\n",
        "\n",
        "As a brief refresher, the KS test finds the maximum difference in the CDFs of two randomly sampled datasets. If the two sets of samples come from the same distribution, the KS statistic will trend to 0 with enough data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "I3XcKxyFTcon"
      },
      "outputs": [],
      "source": [
        "from torcheval.metrics.metric import Metric\n",
        "from scipy.stats import ks_2samp\n",
        "\n",
        "class KS_2Samp(Metric[torch.Tensor]):\n",
        "    def __init__(self, device = None) -> None:\n",
        "        super().__init__(device=device)\n",
        "        # Keep a list of the samples\n",
        "        self._add_state(\"dist_1_samples\", torch.tensor([], device=self.device)) \n",
        "        self._add_state(\"dist_2_samples\", torch.tensor([], device=self.device))\n",
        "\n",
        "    @torch.inference_mode() # turn off autograd and apply some automatic optimizations\n",
        "    def update(self, new_samples_dist_1, new_samples_dist_2):\n",
        "        # When new data comes in, just add them to the list of samples\n",
        "        self.dist_1_samples = torch.cat((self.dist_1_samples, new_samples_dist_1))\n",
        "        self.dist_2_samples = torch.cat((self.dist_2_samples, new_samples_dist_2))\n",
        "        return self\n",
        "\n",
        "    @torch.inference_mode()\n",
        "    def compute(self):\n",
        "        print(\"Computing with\", self.dist_1_samples.shape[0], \"samples\", end=\": \") #just for show\n",
        "        # Let scipy do the hard work\n",
        "        return ks_2samp(\n",
        "            self.dist_1_samples.cpu().detach().numpy(),\n",
        "            self.dist_2_samples.cpu().detach().numpy(),\n",
        "        )\n",
        "\n",
        "    @torch.inference_mode()\n",
        "    def merge_state(self, metrics):\n",
        "        # Merging the states just means concatenating all the samples for each distribution\n",
        "        dist_1_samples = [self.dist_1_samples, ]\n",
        "        dist_2_samples = [self.dist_2_samples, ]\n",
        "        for metric in metrics:\n",
        "            dist_1_samples.append(metric.dist_1_samples)\n",
        "            dist_2_samples.append(metric.dist_2_samples)\n",
        "        self.dist_1_samples = torch.cat(dist_1_samples)\n",
        "        self.dist_2_samples = torch.cat(dist_2_samples)\n",
        "        return self"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9OWEKzk8TpN1"
      },
      "source": [
        "Let's check if the implementation works. We'll sample 10000 elements from a uniform distribution two times and see what the KS statistic is between them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TGoyUZhhTnQI",
        "outputId": "c5130cee-2171-4d59-be99-befd43662317"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing with 10000 samples: KstestResult(statistic=0.0085, pvalue=0.8629266196033867)\n"
          ]
        }
      ],
      "source": [
        "metric = KS_2Samp()\n",
        "metric.update(torch.rand(10000), torch.rand(10000))\n",
        "print(metric.compute()) # statistic should be close to 0 since they have the same underlying distribution"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E_0MdAhJTtBN"
      },
      "source": [
        "Now let's check how the accumulation of state is working by printing the KS statistic with more and more samples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eF5e5JdXTsYC",
        "outputId": "a6e987b9-e322-4923-c9b7-650f8f072989"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing with 500 samples: KstestResult(statistic=0.064, pvalue=0.25760726053364696)\n",
            "Computing with 1000 samples: KstestResult(statistic=0.024, pvalue=0.9357699014782725)\n",
            "Computing with 1500 samples: KstestResult(statistic=0.020666666666666667, pvalue=0.9060256294988647)\n",
            "Computing with 2000 samples: KstestResult(statistic=0.0195, pvalue=0.8416360266077476)\n",
            "Computing with 2500 samples: KstestResult(statistic=0.0136, pvalue=0.9749467010910638)\n",
            "Computing with 3000 samples: KstestResult(statistic=0.02033333333333333, pvalue=0.5646743460194811)\n",
            "Computing with 3500 samples: KstestResult(statistic=0.018571428571428572, pvalue=0.5822115698092439)\n",
            "Computing with 4000 samples: KstestResult(statistic=0.0215, pvalue=0.31359125670612564)\n",
            "Computing with 4500 samples: KstestResult(statistic=0.021333333333333333, pvalue=0.25745106633301273)\n",
            "Computing with 5000 samples: KstestResult(statistic=0.0176, pvalue=0.42096355203665886)\n",
            "\n",
            "=========\n",
            "Resetting Metric\n",
            "=========\n",
            "\n",
            "Computing with 3 samples: KS Statistic for a small random batch: KstestResult(statistic=0.6, pvalue=0.4642857142857143)\n"
          ]
        }
      ],
      "source": [
        "metric = KS_2Samp()\n",
        "for i in range(10):\n",
        "    metric.update(torch.rand(500), torch.rand(500))\n",
        "    print(metric.compute()) # statistic should fall with more data\n",
        "\n",
        "print()\n",
        "print(\"=========\")\n",
        "print(\"Resetting Metric\")\n",
        "print(\"=========\")\n",
        "print()\n",
        "\n",
        "metric.reset()\n",
        "metric.update(torch.rand(3), torch.rand(5))\n",
        "print(\"KS Statistic for a small random batch:\", metric.compute())  # statistic be large with small amounts of data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MxDnqHQ1dkwX"
      },
      "source": [
        "See how the metric accumulates state over time. As we get more examples, the KS statistic gets smaller (roughly). \n",
        "\n",
        "Doing deferred computation of the metric is a useful trick if accumulating state is fast but metric computation is slow. This is true for the KS test, because computing requires sorting of the samples tensors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VW4wOkG2T1OH"
      },
      "source": [
        "# Running In A Distributed Setting\n",
        "\n",
        "Since we've built our metric on top of torcheval, running it in a distributed loop is a breeze"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "wLfQA0DOT5Sd"
      },
      "outputs": [],
      "source": [
        "def distributed_ksloop():\n",
        "    dist.init_process_group(backend=\"gloo\")\n",
        "    metric = KS_2Samp()\n",
        "    for i in range(10):\n",
        "        metric.update(torch.rand(500), torch.rand(500))\n",
        "\n",
        "    return sync_and_compute(metric)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0S-MA5oBT5kX",
        "outputId": "dc474ee0-16a2-43d8-c8e5-3d3ccac0bccd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing with 20000 samples: "
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "KstestResult(statistic=0.029200000000000004, pvalue=7.680638788531949e-08)\n"
          ]
        }
      ],
      "source": [
        "batch_values = pet.elastic_launch(lc, entrypoint=distributed_ksloop)()\n",
        "print(batch_values[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UL4ZjzhldkwX"
      },
      "source": [
        "Notice that our final result is computed with 20,000 samples, 5,000 from each process!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7L4hxlKkT-Wm"
      },
      "source": [
        "# Conclusions\n",
        "\n",
        "Thanks for checking out TorchEval! \n",
        "\n",
        "Please check out our [Docs](https://pytorch.org/torcheval/main/), [More Examples](https://pytorch.org/torcheval/main/metric_example.html), and [Github Repo](https://github.com/pytorch/torcheval). \n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "fileHeader": "",
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "bento_kernel_default"
    },
    "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": "f81b9575fc3922c99b6f4b5377ce97c20846c880e993ba7a7e2470e95e6ea61a"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}