torch-neuronx/inference/tv_pretrained_vgg_inference_on_trn1.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "7e9a2b5b",
            "metadata": {},
            "source": [
                "# Torchvision Pretrained VGG Inference on Trn1 / Inf2"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "2f4b19d5",
            "metadata": {},
            "source": [
                "## Introduction\n",
                "\n",
                "This notebook demonstrates how to compile and run a Torchvision ResNet model for accelerated inference on Neuron. This notebook will use the [`vgg11`](https://pytorch.org/vision/main/models/generated/torchvision.models.vgg11.html) or [`vgg16`](https://pytorch.org/vision/main/models/generated/torchvision.models.vgg16.html) model, which is primarily used for arbitrary image classification tasks.\n",
                "\n",
                "This Jupyter notebook should be run on a Trn1 or Inf2 instance (`trn1.2xlarge` or `inf2.xlarge` or larger)."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Verify that this Jupyter notebook is running the Python kernel environment that was set up according to the [PyTorch Installation Guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/torch-neuronx.html#setup-torch-neuronx). You can select the kernel from the 'Kernel -> Change Kernel' option on the top of this Jupyter notebook page."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "23575b62",
            "metadata": {},
            "source": [
                "## Install Dependencies\n",
                "This tutorial requires the following pip packages:\n",
                "\n",
                "- `torch-neuronx`\n",
                "- `neuronx-cc`\n",
                "- `torchvision`\n",
                "\n",
                "Most of these packages will be installed when configuring your environment using the Trn1 setup guide."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "9dc8c2b1",
            "metadata": {},
            "source": [
                "## Compile the model into an AWS Neuron optimized TorchScript"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5ca0f1b2",
            "metadata": {},
            "source": [
                "In the following section, we load the model, get s sample input, run inference on CPU, compile the model for Neuron using `torch_neuronx.trace()`, and save the optimized model as `TorchScript`.\n",
                "\n",
                "`torch_neuronx.trace()` expects a tensor or tuple of tensor inputs to use for tracing, so we convert the input image into a tensor."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "7ce06089",
            "metadata": {},
            "outputs": [],
            "source": [
                "from PIL import Image\n",
                "import requests\n",
                "\n",
                "import torch\n",
                "import torch_neuronx\n",
                "from torchvision import models\n",
                "from torchvision.transforms import functional\n",
                "\n",
                "# Create the feature extractor and model\n",
                "model = models.vgg11(pretrained=True) # could also choose vgg16\n",
                "model.eval()\n",
                "\n",
                "# Get an example input\n",
                "url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
                "image = Image.open(requests.get(url, stream=True).raw)\n",
                "image = image.convert('RGB')\n",
                "image = functional.resize(image, (224, 224))\n",
                "image = functional.to_tensor(image)\n",
                "image = torch.unsqueeze(image, 0)\n",
                "\n",
                "# Run inference on CPU\n",
                "output_cpu = model(image)\n",
                "\n",
                "# Compile the model\n",
                "model_neuron = torch_neuronx.trace(model, image)\n",
                "\n",
                "# Save the TorchScript for inference deployment\n",
                "filename = 'model.pt'\n",
                "torch.jit.save(model_neuron, filename)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "509f20ef",
            "metadata": {},
            "source": [
                "## Run inference and compare results\n",
                "\n",
                "In this section we load the compiled model, run inference on Neuron, and compare the CPU and Neuron outputs using the ImageNet classes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c1832297",
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "import urllib\n",
                "\n",
                "# Load the TorchScript compiled model\n",
                "model_neuron = torch.jit.load(filename)\n",
                "\n",
                "# Run inference using the Neuron model\n",
                "output_neuron = model_neuron(image)\n",
                "\n",
                "# Compare the results\n",
                "print(f\"CPU tensor:    {output_cpu[0][0:10]}\")\n",
                "print(f\"Neuron tensor: {output_neuron[0][0:10]}\")\n",
                "\n",
                "# Download and read the ImageNet classes\n",
                "urllib.request.urlretrieve(\"https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json\",\"imagenet_class_index.json\")\n",
                "with open(\"imagenet_class_index.json\", \"r\") as file:\n",
                "    class_id = json.load(file)\n",
                "    id2label = [class_id[str(i)][1] for i in range(len(class_id))]\n",
                "\n",
                "# Lookup and print the top-5 labels\n",
                "top5_cpu = output_cpu[0].sort()[1][-5:]\n",
                "top5_neuron = output_neuron[0].sort()[1][-5:]\n",
                "top5_labels_cpu = [id2label[idx] for idx in top5_cpu]\n",
                "top5_labels_neuron = [id2label[idx] for idx in top5_neuron]\n",
                "print(f\"CPU top-5 labels:    {top5_labels_cpu}\")\n",
                "print(f\"Neuron top-5 labels: {top5_labels_neuron}\")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python (Neuron PyTorch)",
            "language": "python",
            "name": "pytorch_venv"
        },
        "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"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}