How to train a fastai model and run it in the browser

This post covers an end-to-end example project of training a resnet model with fastai and PyTorch, exporting it to ONNX and running it in the browser inside a React.js app.

Try the demo!

Beginner-friendly tutorials for training a deep learning model with fast.ai, exporting a PyTorch model to ONNX or creating a frontend web app with React.js are widely available, but when it comes to combining these different steps into a real project, information is more scarce. Hence, the goal of this post (and the accompanying notebook and repository) is to close this gap and show you how the training, export and deployment of a (close to) state-of-the-art deep learning model that runs in any modern web browser with a useable user interface comes together.

Train a resnet with fastai and export to ONNX

This post assumes that you are familiar with training a model using fastai in general. If you are new to fastai, I wholeheartedly recommend their free MOOC. The training process used is described in and can be reproduced with this jupyter notebook. The notebook also includes the export to ONNX, which is a built-in feature of PyTorch.

Run the model in the browser

Deployment of the model to the browser also requires a web app that accepts the input data – in this case an image – and displays the predictions. You can quickly create a react app using create-react-app. To get a somewhat nice looking UI going, this project uses the MATERIAL-UI framework, which provides a lot of useful components. You can find the full source code of this project in the project’s GitHub repository.

Loading and running the model is achieved with onnxjs. Here are some tips when using onnxjs:

Model warm-up

While you load the model, you most likely want to display a loading indicator to the user. I would recommend you also warm-up the model right after loading and keep the loading indicator shown to the user. The warm-up is simply the very first forward pass (i.e. inference) on the loaded model. This will make the inference much faster the second time. So for a smooth user experience, it can make sense to load and warm-up at the start so that running the model is subsequently much faster. Warm-up can be performed using a random input tensor.

Loading the input image tensor

Getting an input image into a tensor format is somewhat tricky in the browser. I use an HTML dropzone, which can handle drag & dropped files as well as a regular file choser dialog. If you are using react, I recommend the react-dropzone component. Once an image file is dropped there, you will need to load the image file to a canvas, which seems to be the only way to get the raw image data from an encoded file like a JPG in JavaScript.

Preprocessing the input tensor

Make sure that you are normalizing the input tensor using the same statistics that you used when training your model. In this case, these were the imagenet_stats.

Interpreting the outputMap

Running the model asynchronously returns an array of numbers, which has the same length as your classifier has classes. Each number thus representens the activation for a given class. To interpret this result, I selected only the top five activations and then computed a percentage “probability”. Note that this isn’t exactly scientific, but gives the user an idea of how confident the model is when making a prediction.

Demo time: Try it out!

Try out the live demo version of the dog breed classifier! It should work in all modern browsers, but depending on the hardware, especially on low-end phones, it might not be able to load or run the model. Be aware that the model itself is about 48MB is size, so you probably want to have a good connection.

The best part: none of the images you give to the classifer ever leave your device! All of the inference is done locally, in your browser.

Feedback welcome

I could and want to say a lot more about the joys and pains of making this little project, but I also want to get this information out there as soon as possible. So if you have any specific questions, suggestions or recommendations, please do not hesitate to contact me on twitter @davidpfahler. I almost certainly made mistakes in the implementation or explanation, so if you find any, please let me know!

I also created a thread on the fastai forums, so if you are active there, please join the discussion!

Acknowledgements

I want to thank first and foremost the amazing team at fast.ai and the community at the fast.ai forums, who are always most helpful. This little pet project obiously stands on the shoulders of (open-source) giants. Thanks to all of them!

davidpfahler

David Pfahler

I am a German trainee lawyer, software engineer and data scientist. I write about these topics here.

Read More