{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch_mlir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet18, ResNet18_Weights\n",
    "model = resnet18(weights=ResNet18_Weights.DEFAULT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "compiled_module = torch_mlir.compile(model, example_args=[torch.ones((1, 3, 224, 224))], output_type=torch_mlir.OutputType.TOSA)\n",
    "\n",
    "from torch_mlir_e2e_test.tosa_backends.linalg_on_tensors import LinalgOnTensorsTosaBackend\n",
    "from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import RefBackendLinalgOnTensorsBackend\n",
    "# backend = LinalgOnTensorsTosaBackend()\n",
    "backend = RefBackendLinalgOnTensorsBackend()\n",
    "compiled_module = torch_mlir.compile(model, example_args=[torch.ones((1, 3, 224, 224))], output_type=torch_mlir.OutputType.LINALG_ON_TENSORS)\n",
    "runnable = backend.compile(compiled_module)\n",
    "jit_module = backend.load(runnable)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "94063946"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outfile = open(\"resnet18.mlir\", \"w\")\n",
    "outfile.write(str(compiled_module))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "# get example image\n",
    "from PIL import Image\n",
    "import requests\n",
    "from io import BytesIO\n",
    "from torchvision import transforms\n",
    "\n",
    "def load_and_preprocess_image(url: str):\n",
    "    headers = {\n",
    "        'User-Agent':\n",
    "        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36'\n",
    "    }\n",
    "    img = Image.open(requests.get(url, headers=headers,\n",
    "                                  stream=True).raw).convert(\"RGB\")\n",
    "    # preprocessing pipeline\n",
    "    preprocess = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                             std=[0.229, 0.224, 0.225]),\n",
    "    ])\n",
    "    img_preprocessed = preprocess(img)\n",
    "    return torch.unsqueeze(img_preprocessed, 0)\n",
    "\n",
    "image_url = \"https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg\"\n",
    "\n",
    "\n",
    "img = load_and_preprocess_image(image_url)\n",
    "\n",
    "print(img.shape)\n",
    "# write img to file\n",
    "torch.save(img.numpy(), \"../examples/test_image.pt\")\n",
    "# write to json \n",
    "import json\n",
    "with open(\"../examples/test_image.json\", \"w\") as f:\n",
    "    json.dump({\n",
    "        \"data\": img.numpy().flatten().tolist(), \n",
    "        \"shape\": list(img.shape)},f)\n",
    "arg = img.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.0.0.dev20230209+cu117'"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16.245480597004644"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import timeit\n",
    "\n",
    "timeit.timeit(lambda: jit_module.forward(arg), number=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.6393603989999974"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "timeit.timeit(lambda: jit_module.forward(arg), number=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.06698099699860904"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "timeit.timeit(lambda: model(img), number=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "18824bd52c2964eef04022de1082fbe6ca8a05a9cab1618bc6c06c0883c4df04"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}