torch-neuronx/inference/hf_pretrained_perceiver_vision_inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7e9a2b5b",
   "metadata": {},
   "source": [
    "# HF Pretrained Perceiver Image Classification Inference on Trn1 / Inf2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2f4b19d5",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook demonstrates how to compile and run a Perceiver vision model for accelerated inference on Neuron. This notebook will use the [`PerceiverForImageClassificationConvProcessing`](https://huggingface.co/deepmind/vision-perceiver-conv) model, which is primarily used for arbitrary image classification tasks. However, the other flavours of the Perceiver vision model can be used similarly (`PerceiverForImageClassificationLearned` and `PerceiverForImageClassificationFourier`)\n",
    "\n",
    "This Jupyter notebook should be run on a Trn1 or Inf2 instance (`trn1.2xlarge` or `inf2.xlarge` or larger)."
   ]
  },
  {
   "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",
    "- `transformers`\n",
    "- `Pillow`\n",
    "\n",
    "Most of these packages will be installed when configuring your environment using the Trn1 setup guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64134be4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
    "!pip install transformers==4.32.0 Pillow==10.0.0"
   ]
  },
  {
   "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": [
    "import torch\n",
    "import torch_neuronx\n",
    "\n",
    "from transformers import AutoImageProcessor\n",
    "from transformers import PerceiverForImageClassificationLearned, PerceiverForImageClassificationFourier, PerceiverForImageClassificationConvProcessing\n",
    "\n",
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "\"\"\"\n",
    "The model classes correspond to the following pretrained models:\n",
    "    PerceiverForImageClassificationLearned -> deepmind/vision-perceiver-learned\n",
    "    PerceiverForImageClassificationFourier -> deepmind/vision-perceiver-fourier\n",
    "    PerceiverForImageClassificationConvProcessing -> deepmind/vision-perceiver-conv\n",
    "\"\"\"\n",
    "\n",
    "# Create the model and image processor\n",
    "model_name = \"deepmind/vision-perceiver-conv\" # Can also be `deepmind/vision-perceiver-conv` or `deepmind/vision-perceiver-fourier`\n",
    "image_processor = AutoImageProcessor.from_pretrained(model_name)\n",
    "model = PerceiverForImageClassificationConvProcessing.from_pretrained(model_name)\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",
    "\n",
    "image_processor = AutoImageProcessor.from_pretrained(\"deepmind/vision-perceiver-conv\")\n",
    "model = PerceiverForImageClassificationConvProcessing.from_pretrained(\"deepmind/vision-perceiver-conv\")\n",
    "\n",
    "inputs = image_processor(images=image, return_tensors=\"pt\").pixel_values\n",
    "\n",
    "# Run inference on CPU\n",
    "output_cpu = model(inputs)['logits'].squeeze(0)\n",
    "\n",
    "# Compile the model\n",
    "model_neuron = torch_neuronx.trace(model, inputs)\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(inputs)['logits'].squeeze(0)\n",
    "\n",
    "# Compare the results\n",
    "print(f\"CPU tensor:    {output_cpu[0:10]}\")\n",
    "print(f\"Neuron tensor: {output_neuron[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.sort(descending=True)[1][:5]\n",
    "top5_neuron = output_neuron.sort(descending=True)[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}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2df2868b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aws_neuron_venv_pytorch",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}