SageMaker Neo

SageMaker Neo is a capability of Amazon SageMaker AI that compiles a trained model for a specific target so it can be trained once and run anywhere — in the cloud or at the edge.

The problem it solves

Getting good inference performance on a given platform normally means hand-tuning for that platform’s instruction set, memory access patterns and input shapes. Most machine learning tooling is specific either to a framework or to a piece of hardware, which leaves manual trial and error. Neo applies the compiler-and-runtime model to that problem.

How it works

Neo has two parts, a compiler and a runtime.

1. The compilation API reads a model exported from its framework.
2. Framework-specific operations are converted into a framework-agnostic intermediate representation.
3. A series of optimisations is applied.
4. Binary code is generated for the optimised operations and written to a shared object library, with the model definition and parameters saved separately.
5. A Neo runtime for the target platform loads and executes the compiled model.

A compilation job is started from the SageMaker AI console, the AWS CLI, a notebook or the SageMaker AI SDK. The compiled model deploys to a SageMaker AI endpoint or to an AWS IoT Greengrass device.

What it supports

Neo optimises Gluon, Keras, MXNet, PyTorch, TensorFlow, TensorFlow Lite and ONNX models for inference on Android, Linux and Windows, across processors from Ambarella, ARM, Intel, NVIDIA, NXP, Qualcomm, Texas Instruments and Xilinx. It targets two families: cloud instances, including AWS Inferentia, and edge devices.

Neo can optimise models whose parameters are in FP32 or quantised to INT8 or FP16.

It is tested against the computer vision models in the framework model zoos, so results are most predictable for that class of model.

When to reach for it

• Deploying to constrained edge hardware where the uncompiled model is too slow or too large
• Reducing the instance size, and so the cost, of a busy real-time endpoint
• Targeting Inferentia without rewriting the model

Compilation is a build-time step, so it fits in a pipeline between training and deployment rather than changing how the model is authored.