Optimize image classification on AWS IoT Greengrass using ONNX Runtime

Introduction

Performing machine learning inference on edge devices using models trained in the cloud has become a popular use case in Internet of Things (IoT) as it brings the benefits of low latency, scalability, and cost savings. When deploying models to edge devices with limited compute and memory, developers have the challenge to manually tune the model to achieve the desired performance. In this blog post, I will discuss an example on how to use the ONNX Runtime on AWS IoT Greengrass to optimize image classification at the edge.

ONNX is an open format built to represent any type of machine learning or deep learning model while making it easier to access hardware optimizations. It provides a standard format for interoperability between different machine learning frameworks. You can train an image classification model using one of your preferred frameworks (TensorFlow, PyTorch, MxNet, and more) and then export it to ONNX format. To maximize performance, you can use your ONNX models with an optimized inference framework, like ONNX Runtime. ONNX Runtime is an open source project designed to accelerate machine learning inference across a variety of frameworks, operating systems, and hardware platforms with a single set of APIs. While this blog post focuses on an example for image classification, you can use ONNX for a wide range of use cases, like object detection, image segmentation, speech and audio processing, machine comprehension and translation, and more.

AWS IoT Greengrass is an open source Internet of Things (IoT) edge runtime and cloud service that helps you build, deploy, and manage IoT applications on your devices. You can use AWS IoT Greengrass to build edge applications using software modules, called components, that can connect your edge devices to AWS or third-party services. There are several AWS-provided machine learning components that can be used to perform inference on remote devices, with locally generated data, using models trained in the cloud. You can also build your custom machine learning components which can be divided in two categories: components for deploying and updating your machine learning models and runtimes at the edge as well as components that contain the necessary application logic for performing machine learning inference.

Solution Overview

In this example, you will learn how to build and deploy a custom component for image classification on AWS IoT Greengrass. The below architecture and steps represent a possible implementation for this solution.

1. Train a model using your preferred framework and export it to ONNX format, or use a pre-trained ONNX model. You can use Amazon SageMaker Studio and Amazon SageMaker Pipelines to automate this process.

In this blog post, you will be using a pre-trained ResNet-50 model in ONNX format for image classification available from the ONNX Model Zoo. ResNet-50 is a convolutional neural network with 50 layers and the pre-trained version of the model can classify images into a thousand object categories, such as keyboard, mouse, pencil, and many animals.

2. Build and publish the necessary AWS IoT Greengrass components:

• An ONNX Runtime component that contains the necessary libraries to run the ONNX model.
• A component for inference that contains the necessary code, the ResNet-50 model in ONNX format as well as some labels and sample images that will be used for classification. This component will have a dependency on the ONNX Runtime component.

3. Deploy the component on the target device. Once the component is running, it will classify the sample images and publish the results back to AWS IoT Core to the topic demo/onnx. AWS IoT Core is a managed AWS service that let’s you connect billions of IoT devices and route trillions of messages to AWS services without managing infrastructure.

Prerequisites

To be able to run through the steps in this blog post, you will need:

• Basic knowledge of AWS IoT Core and AWS IoT Greengrass.
• Basic understanding of machine learning concepts.
• An AWS Account.  If you don’t have an AWS Account, follow the instructions to create one.
• AWS Command Line Interface (AWS CLI) and Git installed.
• An AWS Identity and Access Management (IAM) user with the necessary permissions for creating and managing AWS resources.

Implementation walkthrough

Initial setup

As part of the initial setup for the environment, there are several resources that you need to provision. All the resources need to be provisioned in the same region. This guide is using the eu-central-1 region. Follow the steps below to get started:
1. The component’s artifacts are going to be stored in an Amazon Simple Storage Service (Amazon S3) bucket. To create an Amazon S3 bucket, follow the instructions from the user guide.
2. To emulate a device where we will deploy the component, you will use an AWS Cloud9 environment and then install AWS IoT Greengrass client software. To perform these steps, follow the instructions from the AWS IoT Greengrass v2 workshop, sections 2 and 3.1.
3. On the AWS Cloud9 environment, make sure you have python 3.6.9 as well as pip 23.0 or higher installed.

Build and publish the ONNX Runtime and inference components

In the next section, you will build and publish the custom components by using AWS CLI, either from a terminal on the local machine or in an AWS Cloud9 environment.

To upload the artifacts to the Amazon S3 bucket created as part of the initial setup, follow the next steps:
1. Clone the git repository that contains the component’s artifacts and recipe:

git clone https://github.com/aws-samples/aws-iot-gg-onnx-runtime.git

2. Navigate to the artifacts folder and zip the files:

cd aws-iot-gg-onnx-runtime/artifacts/com.demo.onnx-imageclassification/1.0.0 
zip -r greengrass-onnx.zip .

3. Upload the zip file to the Amazon S3 bucket that you created in the initial setup:

aws s3 cp greengrass-onnx.zip s3://{YOUR-S3-BUCKET}/greengrass-onnx.zip

To publish the components, perform the following steps:
1. Open the recipe file aws-iot-gg-onnx-runtime/recipes/com.demo.onnx-imageclassification-1.0.0.json in a text editor. Below you have the command to navigate to the recipes directory:

cd aws-iot-gg-onnx-runtime/recipes/

2. Replace the Amazon S3 bucket name in artifacts URI with your own bucket name defined above:

"Artifacts": [
    {
      "URI": "s3://{YOUR-S3-BUCKET}/greengrass-onnx.zip",
      "Unarchive": "ZIP"
    }
  ]

3. Before publishing the component, make sure that you are using the same region where you created the resources in the initial setup. You can set your default region by using the following command:

aws configure set default.region eu-central-1

4. Publish the ONNX Runtime component:

aws greengrassv2 create-component-version --inline-recipe fileb://com.demo.onnxruntime-1.0.0.json

5. Publish the component that will perform the image classification and that has a dependency on the ONNX Runtime:

aws greengrassv2 create-component-version --inline-recipe fileb://com.demo.onnx-imageclassification-1.0.0.json

6. To verify that the components were published successfully, navigate to the AWS IoT Console, go to Greengrass Devices >> Components. In the My Components tab, you should see the two components that you just published:

Deploy the component to a target device

1. To deploy the component to a target device, make sure that you have provisioned an AWS Cloud9 environment with AWS IoT Greengrass client software installed.
2. To setup the necessary permissions for the Greengrass device, make sure that the service role associated with the Greengrass device has permissions to retrieve objects from the Amazon S3 bucket you previously created as well as permissions to publish to the AWS IoT topic demo/onnx.
3. To deploy the component to the target device, go to the AWS IoT Console, navigate to Greengrass Devices >> Deployments and choose Create.
4. Fill in the deployment name as well as the name of the core device you want to deploy to.

5. In the Select Components section, select the component com.demo.onnx-imageclassification.
6. Leave all other options as default and choose Next until you reach the Review section of your deployment and then choose Deploy.
7. To monitor the logs and progress of the components’ deployment, you can open the log file of Greengrass core device on the AWS Cloud9 environment with the following command:

sudo tail -f /greengrass/v2/logs/greengrass.log

8. Please note that the ONNX Runtime component, com.demo.onnxruntime, is automatically installed since the image classification component that we selected for deployment has a dependency on it.

Test the ONNX image classification component deployment

When the image classification component is in the running state, it will loop through the files in the images folder and it will classify them. The results are published to AWS IoT Core to the topic demo/onnx.

To understand this process, let’s have a look at some code snippets from the image classification component:
1. To check the sample images so that you can later compare them with the predicted labels, please open the images located in aws-iot-gg-onnx-runtime/artifacts/com.demo.onnx-imageclassification/1.0.0/images folder.
2. The predict function shown below starts an inference session using the ONNX Runtime and the pre-trained ResNet-50 neural network in ONNX format.

def predict(modelPath, labelsPath, image):
    labels = load_labels(labelsPath)
    # Run the model on the backend
    session = onnxruntime.InferenceSession(modelPath, None)

3. The image is initially preprocessed and then passed as an input parameter to the inference session. Please note that ResNet-50 model uses images of 224 x 224 pixels.

image_data = np.array(image).transpose(2, 0, 1)
input_data = preprocess(image_data)
start = time.time()
raw_result = session.run([], {input_name: input_data})
end = time.time()

4.  From the inference result, you extract the label of the image, and you also calculate the inference time in milliseconds.

inference_time = np.round((end - start) * 1000, 2)
idx = np.argmax(postprocess(raw_result))
inferenceResult = {
	"label": labels[idx],
	"inference_time": inference_time
}

5. The image classification component loops through the files present in the images folder and invokes the predict function. The results are published to AWS IoT Core to the demo/onnx topic every 5 seconds.

for img in os.listdir(imagesPath):
        request = PublishToIoTCoreRequest()
        request.topic_name = topic
        image = Image.open(imagesPath + "/" + img)
        pred = predict(modelPath, labelsPath, image)
        request.payload = pred.encode()
        request.qos = qos
        operation = ipc_client.new_publish_to_iot_core()
        operation.activate(request)
        future_response = operation.get_response().result(timeout=5)
        print("successfully published message: ", future_response)
        time.sleep(5)

To test that the results have been published successfully to the topic, go to AWS IoT Console, navigate to MQTT Client section and subscribe to the topic demo/onnx. You should see the inference results like in the screenshot below:

Cleaning up

It is a best practice to delete resources you no longer want to use. To avoid incurring additional costs on your AWS account, perform the following steps:
1. Delete the AWS Cloud9 environment where the AWS IoT Greengrass software was installed:

aws cloud9 delete-environment --environment-id <your environment id>

2. Delete the Greengrass core device:

aws greengrassv2 delete-core-device --core-device-thing-name <thing-name>

3. Delete the Amazon S3 bucket where the artifacts are stored:

aws s3 rb s3://{YOUR-S3-BUCKET} --force

Conclusion

In this blog post, I showed you how you can build and deploy a custom component on AWS IoT Greengrass that uses the ONNX Runtime to classify images. You can customize this component by adding additional images, or by using a different model in ONNX format to make predictions.

To take a deeper dive into AWS IoT Greengrass, including how to build custom components, please check the AWS IoT Greengrass Workshop v2. You can also read the developer guide to get more information on how to customize machine learning components.

About the author