| {"nbformat": 4, "cells": [{"source": "# Importing an ONNX model into MXNet\n\nIn this tutorial we will:\n\n- learn how to load a pre-trained ONNX model file into MXNet.\n- run inference in MXNet.\n\n## Prerequisites\nThis example assumes that the following python packages are installed:\n- [mxnet](http://mxnet.incubator.apache.org/install/index.html)\n- [onnx](https://github.com/onnx/onnx) (follow the install guide)\n- Pillow - A Python Image Processing package and is required for input pre-processing. It can be installed with ```pip install Pillow```.\n- matplotlib", "cell_type": "markdown", "metadata": {}}, {"source": "from PIL import Image\nimport numpy as np\nimport mxnet as mx\nimport mxnet.contrib.onnx as onnx_mxnet\nfrom mxnet.test_utils import download\nfrom matplotlib.pyplot import imshow", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "### Fetching the required files", "cell_type": "markdown", "metadata": {}}, {"source": "img_url = 'https://s3.amazonaws.com/onnx-mxnet/examples/super_res_input.jpg'\ndownload(img_url, 'super_res_input.jpg')\nmodel_url = 'https://s3.amazonaws.com/onnx-mxnet/examples/super_resolution.onnx'\nonnx_model_file = download(model_url, 'super_resolution.onnx')", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "## Loading the model into MXNet\n\nTo completely describe a pre-trained model in MXNet, we need two elements: a symbolic graph, containing the model's network definition, and a binary file containing the model weights. You can import the ONNX model and get the symbol and parameters objects using ``import_model`` API. The paameter object is split into argument parameters and auxilliary parameters.", "cell_type": "markdown", "metadata": {}}, {"source": "sym, arg, aux = onnx_mxnet.import_model(onnx_model_file)", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "We can now visualize the imported model (graphviz needs to be installed)", "cell_type": "markdown", "metadata": {}}, {"source": "mx.viz.plot_network(sym, node_attrs={\"shape\":\"oval\",\"fixedsize\":\"false\"})", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "## Input Pre-processing\n\nWe will transform the previously downloaded input image into an input tensor.", "cell_type": "markdown", "metadata": {}}, {"source": "img = Image.open('super_res_input.jpg').resize((224, 224))\nimg_ycbcr = img.convert(\"YCbCr\")\nimg_y, img_cb, img_cr = img_ycbcr.split()\ntest_image = np.array(img_y)[np.newaxis, np.newaxis, :, :]", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "## Run Inference using MXNet's Module API\n\nWe will use MXNet's Module API to run the inference. For this we will need to create the module, bind it to the input data and assign the loaded weights from the two parameter objects - argument parameters and auxilliary parameters.\n\nTo obtain the input data names we run the following line, which picks all the inputs of the symbol graph excluding the argument and auxiliary parameters:", "cell_type": "markdown", "metadata": {}}, {"source": "data_names = [graph_input for graph_input in sym.list_inputs()\n if graph_input not in arg and graph_input not in aux]\nprint(data_names)", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "```['1']```", "cell_type": "markdown", "metadata": {}}, {"source": "mod = mx.mod.Module(symbol=sym, data_names=data_names, context=mx.cpu(), label_names=None)\nmod.bind(for_training=False, data_shapes=[(data_names[0],test_image.shape)], label_shapes=None)\nmod.set_params(arg_params=arg, aux_params=aux, allow_missing=True, allow_extra=True)", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "Module API's forward method requires batch of data as input. We will prepare the data in that format and feed it to the forward method.", "cell_type": "markdown", "metadata": {}}, {"source": "from collections import namedtuple\nBatch = namedtuple('Batch', ['data'])\n\n# forward on the provided data batch\nmod.forward(Batch([mx.nd.array(test_image)]))", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "To get the output of previous forward computation, you use ``module.get_outputs()`` method.\nIt returns an ``ndarray`` that we convert to a ``numpy`` array and then to Pillow's image format", "cell_type": "markdown", "metadata": {}}, {"source": "output = mod.get_outputs()[0][0][0]\nimg_out_y = Image.fromarray(np.uint8((output.asnumpy().clip(0, 255)), mode='L'))\nresult_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\")\nresult_img.save(\"super_res_output.jpg\")", "cell_type": "code", "execution_count": null, "outputs": [], "metadata": {}}, {"source": "\nYou can now compare the input image and the resulting output image. As you will notice, the model was able to increase the spatial resolution from ``256x256`` to ``672x672``.\n\n```eval_rst\n.. list-table::\n :header-rows: 1\n\n * - Input Image \n - Output Image \n * - .. figure:: https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/doc/tutorials/onnx/images/super_res_input.jpg?raw=true :alt: input input \n - .. figure:: https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/doc/tutorials/onnx/images/super_res_output.jpg?raw=true :alt: output output \n```\n\n<!-- INSERT SOURCE DOWNLOAD BUTTONS -->\n", "cell_type": "markdown", "metadata": {}}], "metadata": {"display_name": "", "name": "", "language": "python"}, "nbformat_minor": 2} |