MAX Engine is a high-performance AI compiler and runtime designed to deliver low latency, and high-throughput inference for AI applications. We've shared how you can get started quickly with MAX in this getting started guide, and how you can deploy MAX Engine optimized models as a microservice using MAX Serving. In this blog post, I’ll show you how MAX Engine optimized model provide huge performance gains, while still delivering highly accurate inference results.

When you provide MAX Engine with a TensorFlow, PyTorch, or ONNX model, it performs several graph-level optimizations such as operator and kernel fusion, kernel specialization, memory layout optimization, shape inference, constant folding and more. These graph-level optimizations do not change the underlying computation in the graph, instead they restructure the graph to perform the operations much faster and more efficiently while maintaining high numerical accuracy.

Through the rest of the blog post, I'll walk you through setting up and run an experiment to compare MAX Engine's inference performance and accuracy against native TensorFlow execution on the famous ImageNet dataset. We'll see that MAX Engine delivers massive performance gains while maintaining the same high accuracy you get from TensorFlow.

MAX Engine accuracy on the ImageNet dataset using ResNet50 model

The ImageNet dataset is one of the most influential dataset of all time, since it set the modern AI revolution in motion with the famous ImageNet Large Scale Visual Recognition Challenge (ILSVRC) in 2012. The training dataset is huge, at 1.2 million images for training 50,000 images for validation. This makes it a good dataset for benchmarking and testing model performance and accuracy.

The ResNet50 model (Kaiming He et. al) won the ILSVRC competition in 2015 and is now a standard for image classification, thanks to it's high accuracy and ease of training. ResNet50 is also commonly used as a backbone architecture for various other models such as the R-CNN family of models and some YOLO family of models used for object detection.

Due to it's popularity, deep learning frameworks like TensorFlow and PyTorch include highly optimized implementations of ResNet50. However, as you'll see in this example, MAX Engine can squeeze out 2.4x faster inference performance over native TensorFlow, while maintaining the same high accuracy. You can see the inference accuracy, performance results below:

‍

In the table, you can see that the validation accuracy is identical for both TensorFlow+Keras model execution and MAX Engine execution for 50,000 images in the ImageNet dataset. At the same time you get 2.4x throughput speed up and sub 100 millisecond inference for a batch size of 8. Checkout performance.modular.com for more performance data.

Download and prepare the dataset

You can get the ImageNet dataset directly from the ImageNet website by signing up and requesting an access key to download the data. The dataset is also available on Kaggle as an alternative source. The entire dataset is approximately 167GB of data, but I’ll only be using the approximately 7GB of validation dataset to measure accuracy. I also used this helpful script from the good people on the TensorFlow team to convert the raw images into TFRecord format which is far more efficient to read and process using TensorFlow’s tf.data API. I used an Amazon EC2 c6i.4xlarge instance running Ubuntu 20.04 to run this example. Make sure you have sufficient storage space for the dataset. All of the code below is available on GitHub.

Setup data loading pipeline

In order to feed our model with batches of images, we have to set up a data loading pipeline. The data loading pipeline includes creating a TFRecordDataset, preprocessing images and performing data augmentation. The code for that is below:

Make sure you update the data_dir = '/path/to/imagenet/dataset/tf-records/validation/*' to where your ImageNet validation TFRecord files are.

Download a ResNet50 model

We can download ResNet50 model using Keras APIs and save it to TensorFlow saved_model format since MAX Engine expects models to be in this format. MAX Engine also supports TorchScript and ONNX model formats.

Inference on the ImageNet validation dataset using TensorFlow and Keras

First, let's run the inference loop over the ImageNet validation dataset using TensorFlow only, and report its performance and accuracy results.

Output

We get about ~41 images/second with native TensorFlow inference, and you can see the full summary below:

Inference on the ImageNet validation dataset using MAX Engine

Now let's take a look at MAX Engine performance and accuracy. To use the same model with MAX Engine you have to first compile the model with MAX Engine and all this takes is 3 lines of code:

After that you just have to replace: model(validation_ds) with model.execute(input_1=validation_ds)

And the rest of the code remains the same (except for any variable name changes you want to make).

Output:

This gives us about 2.4x speedup on this model compared to TensorFlow execution. But we're more interested in accuracy and you can see that the accuracy is identical to our previous TensorFlow only section.

Conclusion

With MAX Engine you get huge performance gains over optimized TensorFlow implementation while maintaining high accuracy! Check out the MAX performance dashboard for speedups MAX can deliver for a range of popular computer vision, natural language, recommender systems and other models: performance.modular.com. The code example used in this blog post is available on GitHub. Download MAX and try it out and share your feedback with us!

Until next time!🔥

Additional resources:

• Get started with downloading MAX
• Download and run MAX examples on GitHub 
• Head over to the MAX docs to learn more about MAX Engine APIs and Mojo programming manual 
• Join our Discord community 
• Contribute to discussions on the Mojo and MAX GitHub

Report feedback, including issues on our Mojo and MAX GitHub tracker