Summary
In this post, we will build a custom image classifier in the cloud using Amazon SageMaker, convert the model to the open ONNX model format, download the ONNX model, then run it in our web browser using ONNX.js.
First, show me a working demo!
If you don't want to train your own model but you still want to see a working demo, you can download an example model and see it in action using this simple single-page-app front end (shared via Glitch):
OK! How do I do this for myself?
If you know a bit of Python, training your own custom image classifier is surprisingly easy using Amazon SageMaker. SageMaker will host a Jupyter Notebook compute environment for you, which you can use to prep your data, train your model, and even deploy your trained model to a fully-managed hosted endpoint (but we won't do that last bit for this blog post, since we want to take a model that we've trained in the cloud and download it for offline inference).
Getting a Jupyter Notebook instance set up
First, log in to your AWS account and go to the Amazon SageMaker web console. If you don't have an AWS account yet, visit https://aws.amazon.com/ to create an account.
Then, look for a 'Create notebook instance' button and click it.
Next, you'll need to pick a name for your notebook instance, and select or create a new IAM role for the notebook instance to assume when it runs. If you're unsure, let the console create a role for you here. You can leave the notebook instance type set to the default of ml.t2.medium. Even though we'll use a relatively low-powered compute instance type here, we will be able to use an on-demand high-powered deep learning optimized instance type for the duration of our model training.
Click 'Create notebook instance' at the bottom of the form to continue.
After a minute or two, your notebook instance will switch status from Pending to Active and you can click Open Jupyter Lab to open the notebook interface.
Next, look on the left-hand sidebar and click the Git icon to clone a repository.
Paste https://github.com/gabehollombe-aws/sagemaker-image-classifier-to-onnx-in-browser.git into the dialog box and click Clone. This will clone a repo containing a sample notebook file for you to use to train your own image classifier.
In the left-hand sidebar, navigate to the cloned repo directory, open the sagemaker directory inside, and open the notebook inside it, named train_and_export_as_onnx.ipynb.
Building a custom image classifier model with Amazon Sagemaker and converting it to ONNX format
Take a look at the train_and_export_as_onnx.ipynb notebook file. You'll see a lot of annotated steps that show how to prep some images for classification and then how to use the AWS SageMaker Python SDK to train a custom image classifier with our image data.
Take note that this notebook will try to use one ml.p3.2xlarge spot instance for the model training which, at the time of this writing, will come out to cost about $0.15 USD in training costs using the sample data. If you don't want to incur any costs for training a model. you can use the pre-trained model linked in the Glitch-hosted app embedded at the top of this post.
In the section titled Grab a bunch of images grouped by folders, one per label class, you'll see we are downloading a collection of images to use for our training.
If you want to use your own custom images instead of these example ones, just modify the notebook to fill the dataset_dir with appropriately named sub-directories (each directory should be named with a label describing the class of images inside it) and put a bunch of example images in each label sub-directory. But, for the purposes of this blog post, I'll assume you're just going to use the set of images from the Caltech 101 data set that the notebook downloads by default.
From the Run menu, select Run All Cells.
It's going to take some time for the notebook to train a custom image classifier model. You'll know it's on the right track because eventually you'll start to see some training log output under the Start the training section.
Eventually (after about 20 minutes or so) the training job will finish. Continue down a bit further in the notebook's output and you should see the cells that download the SageMaker-built model and covert it to the open ONNX format. Find the cell output providing a link to download the ONNX model and click it to get the ONNX model onto your computer.
Finally, we'll need to know a list of all of the label classes that our model will provide scores on on when we use it for inference to classify new inputs. Find the cell showing the list of space-delimited class labels and copy that output to your clipboard for later use.
Using our ONNX image classifer model in the browser with ONNX.js
ONNX.js makes it possible to run inference through ONNX models in the browser (or in Node) and they even have a nice demo website showing how to use ONNX.js with some pre-trained models.
However, I wanted a bit of a nicer interface to play around with, and I wanted to use my own custom image classifier trained via SageMaker, not one of the pre-trained models from the ONNX model zoo. So, I built a little React single-page-app to let you load an ONNX model from your computer into your browser's memory and then perform inferences against images captured from a webcam, image URLs from the Internet, or images that you can drag-and-drop from your computer.
After you've downloaded your custom image classifier ONNX model above, you can use my in-browser inference app to try it out.
Visit https://gabehollombe-aws-sagemaker-image-classifier-to-onnx-in-browser.glitch.me/ to load the inference app in your browser.
Or, check out the repository containing my sample Jupyter notebook and the inference app on GitHub at https://github.com/gabehollombe-aws/sagemaker-image-classifier-to-onnx-in-browser
Awesome! Where can I learn more?
I think you should start with AWS's free Machine Learning training materials.
You can also learn more about building, training, and hosting ML models with SageMaker on the Amazon SageMaker product page.
Top comments (0)