Neural Net in your Phone: From Training to Deployment through ONNX
Current smartphones are powerful enough to run neural networks locally without the need of a cloud server connection. But deploying and running a custom neural network on your phone is not straightforward and it depends on the operating system of your phone. In this post we will go through all the necessary steps to train a custom image classifier network model, export it with ONNX, convert it to Core ML (machine learning framework for iOS apps) and finally deploy it to your iPhone.
Get Images of Mushroom Species from your Region
Mushroom season has already begun in the northern hemisphere. So it would be great to have a mushroom image classifier running locally on our phones in order to identify mushrooms while hiking. For building such image classifier we will need a good training set consisting in dozens of images for each mushroom species.
As an example, we can start with getting the species entity for the iconic mushroom, the fly agaric (Amanita muscaria) which even has its own emoji – 🍄
We can also obtain a thumbnail image to get a better picture of what we are talking about:
In the first argument we will specify the species, and then we will ask for observations with attached photos using “HasImage” option, and observations that have been validated by others using “QualityGrade” option and “Research” property:
We can easily import the images using the “ImageURL” property:
Nice! As a starting point, I’m interested in getting images of the most common mushroom species that can be found in my region (Northern Catalonia) of both toxic and edible mushrooms. Warning: Do not use this classifier for cooking without expert consultation! Toxic mushrooms can be deadly! 💀
Note that we cropped and resized the images in order to have 224x224 pixels:
Now we can do the same for the rest of mushroom species:
Generate the training and test sets:
In order to create the training and test set we need to specify the class labels:
Next we need to import the images and create the “examples” as follows:
We have a total of 850 images; fifty for each of the seventeen mushroom species:
Let’s check the first example:
Nice! Finally, we only need to create randomized training and testing sets:
Perform Net Surgery for Transfer Learning
Starting from a pre-trained model we can make use of net surgery functions to create our own custom mushroom image classification network.
We can check wether the size of the network is reasonably small for current smartphones. As a rule of thumb it shouldn’t surpass 100MB:
Finally we train the net, keeping the pre-trained weights fixed:
We can check the resulting net performance by plotting the confusion matrix plot for the test set (using either NetMeasurements or ClassifierMeasurements ):
We can test the classifier using a photo from an iNaturalist user observation:
It’s a good practice to save our trained model. So in case we re-start the session we won’t need to re-train the net:
Export the Trained Net Model with ONNX
Checking that the conversion to ONNX was successful
A quick way to make sure that the conversion to ONNX was done correctly is to import back the model and test its performance against the original model:
ONNX net models don’t store the encoder and the decoder. So, we need to dress our imported ONNX net with the original encoder and decoder in order to run the model:
Using the dressed net we can compare both net models and check wether the output probabilities are the same:
Great! The roundtrip is working well and now we are ready for converting the ONNX model into Core ML.
Convert the Model to Core ML
In this section we will make extensive use of a Python package called coremltools which Apple freely provides in order to convert external neural net models to Core ML (Apple’s machine learning framework for iOS apps).
Configure Python for ExternalEvaluate
In order to configure your system to evaluate external Python code I recommend you to follow this workflow: https://reference.wolfram.com/language/workflow/ConfigurePythonForExternalEvaluate.html
Setup External Evaluator for running Python code
Once Python has been configured for ExternalEvaluate we need to register it as an external evaluator and start an external session:
Install required packages
For converting an ONNX model into Core ML we need to install two extra packagesusing the terminal:
Specifying the encoder and decoder during the conversion
Core ML models (.mlmodel) work similarly to Wolfram Language models. They also include an encoder and a decoder for the model. So, while converting the ONNX model to Core ML we need to specify the image encoder (preprocessing arguments) and the decoder (the class labels).
Preprocessing arguments (color bias and mean image)
If we click on the “Input” port from our original Wolfram Language model we will see the following panel:
During the conversion we will need to specify that the input type is an “Image” and include the mean image values for each color channel as biases. Furthermore, we will need to specify an image re-scale factor since in the original model pixel values range from 0 to 1 and in Core ML the scale is from 0 to 255.
Save class labels as a text file
Coremltools allow us to specify the class labels of the model using a text file containing each class label in a new line. It is straightforward to export such a text file using Export and StringRiffle:
Python code to convert the ONNX model
The following python code consists in three parts. Importing coremltools package and specifying the path to ONNX model. The code for converting the model and finally the code for saving the resulting Core ML model:
Deploy the Model to the iPhone
Before deploying the model we need to sign the development team of the app:
Once we drop/upload the model inside the Xcode project, we will see the following window for the model. The “Preview” section allow us to test the model directly using Xcode:
Finally, we need to replace the “Mobilenet” Core ML classifier model for our “MushroomsWolframNet” model in the “ImageClassificationViewController” swift file and then press the “Build and then run” button on the top left:
Test the Mushroom Classifier App in the Wild
We did it! We just need to go hiking into a nearby forest and hope to find some mushrooms. 😋
See a few examples from last days (all correctly identified):