{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# For tips on running notebooks in Google Colab, see\n# https://codelin.vip/beginner/colab\n%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "(optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime\n===================================================================================\n\n```{=html}\n
NOTE:
\n```\n```{=html}\n
\n```\n```{=html}\n

As of PyTorch 2.1, there are two versions of ONNX Exporter.

\n```\n```{=html}\n
\n```\nIn this tutorial, we describe how to convert a model defined in PyTorch\ninto the ONNX format using the TorchScript `torch.onnx.export` ONNX\nexporter.\n\nThe exported model will be executed with ONNX Runtime. ONNX Runtime is a\nperformance-focused engine for ONNX models, which inferences efficiently\nacross multiple platforms and hardware (Windows, Linux, and Mac and on\nboth CPUs and GPUs). ONNX Runtime has proved to considerably increase\nperformance over multiple models as explained\n[here](https://cloudblogs.microsoft.com/opensource/2019/05/22/onnx-runtime-machine-learning-inferencing-0-4-release)\n\nFor this tutorial, you will need to install\n[ONNX](https://github.com/onnx/onnx) and [ONNX\nRuntime](https://github.com/microsoft/onnxruntime). You can get binary\nbuilds of ONNX and ONNX Runtime with\n\n``` {.bash}\n%%bash\npip install onnx onnxruntime\n```\n\nONNX Runtime recommends using the latest stable runtime for PyTorch.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Some standard imports\nimport numpy as np\n\nfrom torch import nn\nimport torch.utils.model_zoo as model_zoo\nimport torch.onnx" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Super-resolution is a way of increasing the resolution of images, videos\nand is widely used in image processing or video editing. For this\ntutorial, we will use a small super-resolution model.\n\nFirst, let\\'s create a `SuperResolution` model in PyTorch. This model\nuses the efficient sub-pixel convolution layer described in [\\\"Real-Time\nSingle Image and Video Super-Resolution Using an Efficient Sub-Pixel\nConvolutional Neural Network\\\" - Shi et\nal](https://arxiv.org/abs/1609.05158) for increasing the resolution of\nan image by an upscale factor. The model expects the Y component of the\n`YCbCr` of an image as an input, and outputs the upscaled Y component in\nsuper resolution.\n\n[The\nmodel](https://github.com/pytorch/examples/blob/master/super_resolution/model.py)\ncomes directly from PyTorch\\'s examples without modification:\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Super Resolution model definition in PyTorch\nimport torch.nn as nn\nimport torch.nn.init as init\n\n\nclass SuperResolutionNet(nn.Module):\n def __init__(self, upscale_factor, inplace=False):\n super(SuperResolutionNet, self).__init__()\n\n self.relu = nn.ReLU(inplace=inplace)\n self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))\n self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))\n self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))\n self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))\n self.pixel_shuffle = nn.PixelShuffle(upscale_factor)\n\n self._initialize_weights()\n\n def forward(self, x):\n x = self.relu(self.conv1(x))\n x = self.relu(self.conv2(x))\n x = self.relu(self.conv3(x))\n x = self.pixel_shuffle(self.conv4(x))\n return x\n\n def _initialize_weights(self):\n init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))\n init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))\n init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))\n init.orthogonal_(self.conv4.weight)\n\n# Create the super-resolution model by using the above model definition.\ntorch_model = SuperResolutionNet(upscale_factor=3)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Ordinarily, you would now train this model; however, for this tutorial,\nwe will instead download some pretrained weights. Note that this model\nwas not trained fully for good accuracy and is used here for\ndemonstration purposes only.\n\nIt is important to call `torch_model.eval()` or\n`torch_model.train(False)` before exporting the model, to turn the model\nto inference mode. This is required since operators like dropout or\nbatchnorm behave differently in inference and training mode.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Load pretrained model weights\nmodel_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'\nbatch_size = 64 # just a random number\n\n# Initialize model with the pretrained weights\nmap_location = lambda storage, loc: storage\nif torch.cuda.is_available():\n map_location = None\ntorch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))\n\n# set the model to inference mode\ntorch_model.eval()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Exporting a model in PyTorch works via tracing or scripting. This\ntutorial will use as an example a model exported by tracing. To export a\nmodel, we call the `torch.onnx.export()` function. This will execute the\nmodel, recording a trace of what operators are used to compute the\noutputs. Because `export` runs the model, we need to provide an input\ntensor `x`. The values in this can be random as long as it is the right\ntype and size. Note that the input size will be fixed in the exported\nONNX graph for all the input\\'s dimensions, unless specified as a\ndynamic axes. In this example we export the model with an input of\nbatch\\_size 1, but then specify the first dimension as dynamic in the\n`dynamic_axes` parameter in `torch.onnx.export()`. The exported model\nwill thus accept inputs of size \\[batch\\_size, 1, 224, 224\\] where\nbatch\\_size can be variable.\n\nTo learn more details about PyTorch\\'s export interface, check out the\n[torch.onnx documentation](https://pytorch.org/docs/master/onnx.html).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Input to the model\nx = torch.randn(batch_size, 1, 224, 224, requires_grad=True)\ntorch_out = torch_model(x)\n\n# Export the model\ntorch.onnx.export(torch_model, # model being run\n x, # model input (or a tuple for multiple inputs)\n \"super_resolution.onnx\", # where to save the model (can be a file or file-like object)\n export_params=True, # store the trained parameter weights inside the model file\n opset_version=10, # the ONNX version to export the model to\n do_constant_folding=True, # whether to execute constant folding for optimization\n input_names = ['input'], # the model's input names\n output_names = ['output'], # the model's output names\n dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes\n 'output' : {0 : 'batch_size'}})" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We also computed `torch_out`, the output after of the model, which we\nwill use to verify that the model we exported computes the same values\nwhen run in ONNX Runtime.\n\nBut before verifying the model\\'s output with ONNX Runtime, we will\ncheck the ONNX model with ONNX API. First,\n`onnx.load(\"super_resolution.onnx\")` will load the saved model and will\noutput a `onnx.ModelProto` structure (a top-level file/container format\nfor bundling a ML model. For more information [onnx.proto\ndocumentation](https://github.com/onnx/onnx/blob/master/onnx/onnx.proto).).\nThen, `onnx.checker.check_model(onnx_model)` will verify the model\\'s\nstructure and confirm that the model has a valid schema. The validity of\nthe ONNX graph is verified by checking the model\\'s version, the\ngraph\\'s structure, as well as the nodes and their inputs and outputs.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import onnx\n\nonnx_model = onnx.load(\"super_resolution.onnx\")\nonnx.checker.check_model(onnx_model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let\\'s compute the output using ONNX Runtime\\'s Python APIs. This\npart can normally be done in a separate process or on another machine,\nbut we will continue in the same process so that we can verify that ONNX\nRuntime and PyTorch are computing the same value for the network.\n\nIn order to run the model with ONNX Runtime, we need to create an\ninference session for the model with the chosen configuration parameters\n(here we use the default config). Once the session is created, we\nevaluate the model using the run() API. The output of this call is a\nlist containing the outputs of the model computed by ONNX Runtime.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import onnxruntime\n\nort_session = onnxruntime.InferenceSession(\"super_resolution.onnx\", providers=[\"CPUExecutionProvider\"])\n\ndef to_numpy(tensor):\n return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n\n# compute ONNX Runtime output prediction\nort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}\nort_outs = ort_session.run(None, ort_inputs)\n\n# compare ONNX Runtime and PyTorch results\nnp.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)\n\nprint(\"Exported model has been tested with ONNXRuntime, and the result looks good!\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We should see that the output of PyTorch and ONNX Runtime runs match\nnumerically with the given precision (`rtol=1e-03` and `atol=1e-05`). As\na side-note, if they do not match then there is an issue in the ONNX\nexporter, so please contact us in that case.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Timing Comparison Between Models\n================================\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since ONNX models optimize for inference speed, running the same data on\nan ONNX model instead of a native pytorch model should result in an\nimprovement of up to 2x. Improvement is more pronounced with higher\nbatch sizes.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import time\n\nx = torch.randn(batch_size, 1, 224, 224, requires_grad=True)\n\nstart = time.time()\ntorch_out = torch_model(x)\nend = time.time()\nprint(f\"Inference of Pytorch model used {end - start} seconds\")\n\nort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}\nstart = time.time()\nort_outs = ort_session.run(None, ort_inputs)\nend = time.time()\nprint(f\"Inference of ONNX model used {end - start} seconds\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Running the model on an image using ONNX Runtime\n================================================\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "So far we have exported a model from PyTorch and shown how to load it\nand run it in ONNX Runtime with a dummy tensor as an input.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "For this tutorial, we will use a famous cat image used widely which\nlooks like below\n\n![](https://pytorch.org/tutorials/_static/img/cat_224x224.jpg)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "First, let\\'s load the image, preprocess it using standard PIL python\nlibrary. Note that this preprocessing is the standard practice of\nprocessing data for training/testing neural networks.\n\nWe first resize the image to fit the size of the model\\'s input\n(224x224). Then we split the image into its Y, Cb, and Cr components.\nThese components represent a grayscale image (Y), and the\nblue-difference (Cb) and red-difference (Cr) chroma components. The Y\ncomponent being more sensitive to the human eye, we are interested in\nthis component which we will be transforming. After extracting the Y\ncomponent, we convert it to a tensor which will be the input of our\nmodel.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from PIL import Image\nimport torchvision.transforms as transforms\n\nimg = Image.open(\"./_static/img/cat.jpg\")\n\nresize = transforms.Resize([224, 224])\nimg = resize(img)\n\nimg_ycbcr = img.convert('YCbCr')\nimg_y, img_cb, img_cr = img_ycbcr.split()\n\nto_tensor = transforms.ToTensor()\nimg_y = to_tensor(img_y)\nimg_y.unsqueeze_(0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now, as a next step, let\\'s take the tensor representing the grayscale\nresized cat image and run the super-resolution model in ONNX Runtime as\nexplained previously.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img_y)}\nort_outs = ort_session.run(None, ort_inputs)\nimg_out_y = ort_outs[0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "At this point, the output of the model is a tensor. Now, we\\'ll process\nthe output of the model to construct back the final output image from\nthe output tensor, and save the image. The post-processing steps have\nbeen adopted from PyTorch implementation of super-resolution model\n[here](https://github.com/pytorch/examples/blob/master/super_resolution/super_resolve.py).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')\n\n# get the output image follow post-processing step from PyTorch implementation\nfinal_img = Image.merge(\n \"YCbCr\", [\n img_out_y,\n img_cb.resize(img_out_y.size, Image.BICUBIC),\n img_cr.resize(img_out_y.size, Image.BICUBIC),\n ]).convert(\"RGB\")\n\n# Save the image, we will compare this with the output image from mobile device\nfinal_img.save(\"./_static/img/cat_superres_with_ort.jpg\")\n\n# Save resized original image (without super-resolution)\nimg = transforms.Resize([img_out_y.size[0], img_out_y.size[1]])(img)\nimg.save(\"cat_resized.jpg\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Here is the comparison between the two images:\n\n![](https://pytorch.org/tutorials/_static/img/cat_resized.jpg)\n\nLow-resolution image\n\n![](https://pytorch.org/tutorials/_static/img/cat_superres_with_ort.jpg)\n\nImage after super-resolution\n\nONNX Runtime being a cross platform engine, you can run it across\nmultiple platforms and on both CPUs and GPUs.\n\nONNX Runtime can also be deployed to the cloud for model inferencing\nusing Azure Machine Learning Services. More information\n[here](https://docs.microsoft.com/en-us/azure/machine-learning/service/concept-onnx).\n\nMore information about ONNX Runtime\\'s performance\n[here](https://onnxruntime.ai/docs/performance).\n\nFor more information about ONNX Runtime\n[here](https://github.com/microsoft/onnxruntime).\n" ] } ], "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.12" } }, "nbformat": 4, "nbformat_minor": 0 }