Compile pre-trained models using Amazon SageMaker Neo

This is a simple tutorial to exploit the Amazon SageMaker Neo capability of automatically optimizing machine learning models for inference based on the target platform or device.

Image credits : https://enterpriseiotinsights.com/

What we have : A computer vision pre-trained model from the model zoo. What we want : Compiled/optimized model for inference on a specific platform/device which can later be executed using Neo Deep Learning Runtime.

Here, we use yolo3(You only look once) pre-trained model with darknet base from the Gluon-cv model zoo and optimize it for inference on Raspberry pi device. However, the same process can be applied to any of the pre-trained computer vision models from the model zoo for inference on any target device.

Prerequisites

Python3 installed on the machine along with mxnet and gluoncv libraries. Both mxnet and gluoncv can be installed using the pip command.

pip install --upgrade mxnet gluoncv

By following the three steps as mentioned below, a pre-trained model can be compiled using SageMaker Neo.

Step 1: Export the pre-trained model.

Create a python file ‘export_model.py’ with the following contents and execute it. This exports the pre-trained model by generating a JSON file that includes computational graph and params file that includes pre-trained weights of the model.

import gluoncv as gcv
from gluoncv.utils import export_block
net = gcv.model_zoo.get_model('yolo3_darknet_voc', pretrained=True)
export_block('yolo3_darknet_voc', net, preprocess=True,layout='HWC')

get_model downloads the specified pre-trained model to ~/.mxnet/models directory when pre-trained is set to True(line 3).

export_block generates two files yolo3_darknet53_voc-0000.params and yolo3_darknet53_voc-symbol.json in the same directory as the file export_model.py (line 4).

Step 2: Compress the .params and .json file into a tar.gz file and upload it to an s3 bucket so that the Neo compiler can access it as an input. Let the path of the created s3 bucket with the tar file be,

s3://compile-me/yolo3-darknet-voc.tar.gz

Step 3: Create a compilation job from the Amazon SageMaker console. From your AWS services in your account, navigate to Amazon SageMaker. Click on create compilation job as mentioned in the following picture.

Amazon SageMaker Console, credits : AWS SageMaker docs

In the job settings, choose

1. Job name — Give any random name that you want but make sure that it different from the existing names of the already compiled jobs.
2. IAM role — You can choose an existing role or create a new one in order for you to allow the SageMaker to access other resources. As we have our pre-trained model in an S3 bucket, make sure that role has permission to access that s3 bucket.

In the Input configuration, choose

1. Location of model artifacts: s3 path of the bucket -

s3://compile-me/yolo3-darknet-voc.tar.gz

2. Data input configuration: Trained data shape matrix. Gluoncv yolo3 is trained with colored images of shape 512 x 512 and since the layout used is HWC, input configuration takes the following format

{“data” : [1, 512, 512, 3]}

3. Machine learning framework: The framework using which the model is pre-trained — Mxnet

In the Output configuration, choose

1. Target device: Choose target device or target platform of your choice — rasp3b
2. S3 Output location: We need to the specify the path of an s3 bucket where we want our compiled artifacts to be. You can create an s3 bucket if there are not any and say the path of it be — s3://compiled-artifacts/

After you provide all the required parameters to start the compilation job, submit it.

When we find that the compilation succeeds, go to the output s3 bucket location to find the compiled artifacts 🎉

These compiled artifacts can now be loaded and executed using Neo Deep Learning Runtime Api.

References

Installation - gluoncv 0.9.0 documentation

Compile and Deploy Models with Amazon SageMaker Neo

1. Export trained GluonCV network to JSON - gluoncv 0.9.0 documentation