torch-neuronx/inference/pretrained_unet_inference_on_trn1.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "7e9a2b5b",
            "metadata": {},
            "source": [
                "# Pretrained UNet Inference on Trn1 / Inf2"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "2f4b19d5",
            "metadata": {},
            "source": [
                "## Introduction\n",
                "\n",
                "This notebook demonstrates how to compile and run a UNet model for accelerated inference on Neuron. This notebook will use the [`UNet`](https://github.com/milesial/Pytorch-UNet) model, which is primarily used for arbitrary image semantic segmentation 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",
                "# load the model\n",
                "model = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=False)\n",
                "# load the weights\n",
                "state_dict = torch.hub.load_state_dict_from_url('https://github.com/milesial/Pytorch-UNet/releases/download/v3.0/unet_carvana_scale0.5_epoch2.pth', map_location=\"cpu\")\n",
                "model.load_state_dict(state_dict)\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]}\")"
            ]
        }
    ],
    "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
}