Deploying a PyTorch model on Android
In this section, we will create an Android app that allows you to capture an image using the phone camera and make a prediction (image classification) on the captured image. In Chapter 1 of this book, we trained a Modified National Institute of Standards and Technology (MNIST) model to classify handwritten digits. In Chapter 13, we used tracing to convert the trained MNIST model from the original PyTorch format into an intermediate representation (IR). For our Android app, we will first optimize this traced MNIST model using PyTorch Mobile and then use the optimized model to make predictions (handwritten digit classification) on the captured image. All code for this section is available on GitHub [2].
Converting the PyTorch model to a mobile-friendly format
PyTorch Mobile provides a function, optimize_for_mobile
, that converts a traced PyTorch model object into a mobile-friendly lightweight format. This can be done with the following lines...