Export FastAI ResNet models to ONNX
Tags: machinelearning, fastai, pytorch, onnx
Published: 2021-07-07 23:25:55

> A short guide on FastAI vision model conversion to ONNX. Code included. 👀

--------

## What is FastAI?

FastAI is *"making neural nets uncool again"*. It offers a high-level API to PyTorch. It could be considered the Keras of PyTorch (?). FastAI and the accompanying course taught by Jeremy Howard and Rachel Thomas has a practical approach to deep learning. It encourages students to train DL models from the first minute. BUT I am sure you already have hands-on experience with this framework if you are looking to convert your FastAI models to ONNX.

## What is ONNX?

Open Neural Network Exchange or ONNX is a unified format for deep learning and traditional machine learning models. The idea behind ONNX is to create a common interface for all ML frameworks and increase the interoperability between frameworks and devices. 

> ONNX is an open specification that consists of a definition of an extensible computation graph model, definition of standard data types, and definition of built-in operators. Extensible computation graph and definition of standard data types make up the Intermediate Representation (IR). 

Source [[link](https://github.com/onnx/onnx/blob/master/docs/IR.md)]

![image.png](https://cdn.hashnode.com/res/hashnode/image/upload/v1625693023787/yhXdGzeXJ.png)

Source [[link](https://microsoft.github.io/ai-at-edge/docs/onnx/)]

ONNX, and its implementation - ONNX Runtime, make it easier to put your models into production. You can train your models using the framework of your choice and deploy to a target that uses ONNX Runtime. This way bloated environments with large number of dependencies can be minimized to (pretty much) only ONNX Runtime. There's a growing support to ONNX, and exports are being natively supported by frameworks like PyTorch, MxNet, etc. Find all of them here [[link](https://onnx.ai/supported-tools)]. Although in some cases it might be tricky to export/import, due opset compatibility.

### Why use ONNX and ONNX Runtime?

Couple of reasons here:

- Faster inference [[link](https://medium.com/microsoftazure/faster-and-smaller-quantized-nlp-with-hugging-face-and-onnx-runtime-ec5525473bb7)], [[link](https://cloudblogs.microsoft.com/opensource/2020/12/17/accelerate-simplify-scikit-learn-model-inference-onnx-runtime/)]
- Lower number of dependencies*
- Smaller environment size*
- One, universal target framework for deployment

*See conda environment and dependency comparison at the end of the article.

********


## The Process

FastAI currently doesn't natively support ONNX exports from FastAI learners. But by design FastAI is a high-level API of PyTorch. This allows us to extract the wrapped PyTorch model. And luckily, PyTorch models can be natively exported to ONNX. It's a 2-step process with a couple of gotchas. This guide intends to make it a smooth experience.


![image.png](https://cdn.hashnode.com/res/hashnode/image/upload/v1625693254745/cVtex5wkQ.png)

You can find the entire process in my repository [[link](https://github.com/tkeyo/fastai-onnx)]. It also includes an optional ResNet model training. You can skip it and proceed with model export to ONNX. I included a link to a pre-trained model in the notebooks. _**Or BYOFM - Bring Your Own FastAI Model**_.

## Export (Extract) the PyTorch model

![image.png](https://cdn.hashnode.com/res/hashnode/image/upload/v1625698022279/qH5_G0jep.png)

Let's break down what's happening.

If you check the associated notebooks you will find that I exported the FastAI ResNet learner in the previous steps. And named it *`hot_dog_model_resnet18_256_256.pkl`.* With `load_learner()` I am loading the previously exported FastAI model on **line 7**. If you trained your own model you can skip the load step. Your model is already stored in `learn`.

To get the PyTorch model from the FastAI wrapper we use `model` attribute on `learn` - see **line 12**. I don't want to train the model in subsequent steps thus  I am also setting it to evaluation mode with `eval()`. For more details on `eval()` and `torch.no_grad()` see the discussion [[link](https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615)].

FastAI wraps the PyTorch model with additional layer for convenience - **Softmax**, **Normalization**, and other transformation(defined in FastAI DataBlock API). When using the bare-bones PyTorch model I have to make up for this. Otherwise I'll be getting *weird* results.

First, I define the softmax layer. This will turn inference results into more human readable format. It turn the inference results from something likes this `('not_hot_dog', array([[-3.0275817, 1.2424631]], dtype=float32))` into  `('not_hot_dog', array([[0.01378838, 0.98621166]], dtype=float32))`. Notice the range of inference results - with the added **softmax** layer the results are scaled between **0-1**.

On **line 18**, the normalization layer is defined. I am reusing the suggested ImageNet mean and standard deviation values as described here [[link](https://pytorch.org/vision/stable/models.html)]. If you are interested in an in-depth conversation on the topic of normalization. See this [[link](https://discuss.pytorch.org/t/understanding-transform-normalize/21730/8)].

{% gist https://gist.github.com/tkeyo/ddb5fce50704d1b5a7fc622e6c713494 %}

On **lines 21-25**, I am pulling all together into the final model. This final model will be used for ONNX conversion. FastAI learner also handles resizing but for PyTorch and ONNX this will be handled outside of the model by an extra function.

## Export PyTorch to ONNX

![image.png](https://cdn.hashnode.com/res/hashnode/image/upload/v1625697979846/tgc5n0j7C.png)

PyTorch natively support ONNX exports, I only need to define the export parameters. As you can see we are (re)using the `final_model` for export. On line 5  I am creating a dummy tensor that is used to define the input dimensions of my ONNX model. These dimensions are defined as `batch x channels x height x width - BCHW` format. My FastAI model was trained on images with 256 x 256 dimension which was defined in our FastAI DataBlock API. The same dimensions must be used for the ONNX export - `torch.randn(1, 3, 256, 256)`.

I got this wrong a couple of times - the dummy tensor had different dimensions than the images the model was trained on. *Example: Dummy tensor `torch.randn(1, 3, 320, 320)` while training image dimensions were `3 x 224 x 224`. It took me a while to figure out why I got poor results from my ONNX models.*

{% gist https://gist.github.com/tkeyo/021cd1aebde41f52917b2acc64b2ea59 %}

The `export_param` argument, if set to `True`, includes the parameters of the trained model in the export. It's important to use `True` in this case. We want our model with parameters. As you might have guessed, `export_params=False` exports a model without parameters. Full `torch.onnx` documentation [[link](https://pytorch.org/docs/master/onnx.html)].


## Inference with ONNX Runtime

![onnx_runtime_logo.png](https://cdn.hashnode.com/res/hashnode/image/upload/v1625698218983/_ZorSiRJJ.png)

On **line 10**, I am creating an ONNX runtime inference session and loading the exported model. For debugging purposes, or if you get your hands on an ONNX model with unknown input dimensions. You can run `get_inputs()[0].shape` on the inference session instance to get the expected inputs. If you prefer a GUI, Netron [[link](https://netron.app)] can help you to visualize the architecture of the neural networks.

{% gist https://gist.github.com/tkeyo/6777893acbed14386c455d441714ba84 %}

The inference itself is done by using the `run()` method which returns a numpy array with softmaxed probabilities. See **line 21**.

## Storage, Dependencies & Inference Speed

#### Storage

The advantage of using ONNX Runtime is the small storage footprint compared to PyTorch and FastAI. A conda environment with ONNX Runtime (+ Pillow for convenience) is **~ 25% of the PyTorch** environment and only **~ 15% of the FastAI** environment. Important for serverless deployments.

{% gist https://gist.github.com/tkeyo/4a5c540dda46a40e1cb9b61fdb2aacd6 %}

#### Dependencies

See for yourself.

{% gist https://gist.github.com/tkeyo/ba41848a604af64881828c41c3368f9a %}


#### Inference speed

I mentioned inference speed as an advantage of ONNX. I tested the inference speed of all three versions of the same model. There were negligible differences between inference speed. Other experiments had more favourable results for ONNX. See references in **Why use ONNX and ONNX Runtime?**


## Summary

FastAI is a great tool to get you up and running with model training in a (VERY) short time. It has everything you need to get top notch results with minimal effort in a practical manner. But when it comes to deployment, tools like ONNX & ONNX Runtime can save resource with their smaller footprint and efficient implementation. Hope this guide was helpful and that you managed successfully convert your model to ONNX.

## Repository/Code

[FastAI-ONNX GitHub](https://github.com/tkeyo/fastai-onnx)

------


**Feel free to reach out.👏**

>
>
>
>