A first look at AWS Inferentia
Tags: aws, machinelearning
Published: 2019-12-18 21:52:30

---
title: A first look at AWS Inferentia
published: true
date: 2019-12-14 20:25:08 UTC
tags: aws,machinelearning
canonical_url: https://medium.com/@julsimon/a-first-look-at-aws-inferentia-b9672e8f8b8f
---

Launched at AWS re:Invent 2019, [AWS Inferentia](https://aws.amazon.com/machine-learning/inferentia/) is a high performance machine learning inference chip, custom designed by AWS: its purpose is to deliver cost effective, low latency predictions at scale. Inferentia is present in [Amazon EC2 inf1](https://aws.amazon.com/ec2/instance-types/inf1/) instances, a new family of instances also [launched](https://aws.amazon.com/blogs/aws/amazon-ec2-update-inf1-instances-with-aws-inferentia-chips-for-high-performance-cost-effective-inferencing/) at re:Invent.

In this post, I’d like to show you how to get started with Inferentia and [TensorFlow](https://www.tensorflow.org). Please note that [Apache MXNet](https://mxnet.apache.org), [PyTorch](https://pytorch.org) and [ONNX](https://onnx.ai) are also supported.

#### A primer on Inferentia

The [CMP324](https://www.youtube.com/watch?v=17r1EapAxpk) breakout session is a great introduction to Inferentia, and the Alexa use case is a rare look under the hood. It’s well worth your time.

{% youtube 17r1EapAxpk %}

In a nutshell, each Inferentia chip hosts 4 Neuron Cores. Each one of these implements a “high performance [systolic array](https://en.wikipedia.org/wiki/Systolic_array) matrix multipy engine” (nicely put, Gadi), and is also equipped with a large on-chip cache.

![](https://cdn-images-1.medium.com/max/1024/1*IQmXRvF9uIyNlzb-voaIUQ.png)

NeuronCores are interconnected, which makes it possible to:

- Partition a model across multiple cores (and Inferentia chips, if several are available), storing it 100% in on-cache memory.
- Stream data at full speed through the pipeline of cores, without having to deal with latency caused by external memory access.

![](https://cdn-images-1.medium.com/max/1024/1*xSIJzmWfkef7cxnBSMpdyQ.png)

Alternatively, you can run inference with different models on the same Inferentia chip. This is achieved by partitioning NeuronCores into NeuronCore Groups, and by loading different models on different groups.

#### The Neuron SDK

In order to run on Inferentia, models first need to be compiled to a hardware-optimized representation. Then, they may be loaded, executed and profiled using a specific runtime. These operations can be performed through command-line tools available in the [AWS Neuron SDK](https://github.com/aws/aws-neuron-sdk), or through framework APIs.

[aws/aws-neuron-sdk](https://github.com/aws/aws-neuron-sdk)

Let’s get started!

#### Launching an EC2 instance for model compilation

This first step doesn’t require an inf1 instance. In fact, you should use a **compute-optimized instance** for fast and cost effective compilation. In order to avoid any software configuration, you should also use the [**Deep Learning AMI**](https://aws.amazon.com/machine-learning/amis/), which comes preinstalled with the Neuron SDK and with updated frameworks.

At the time of writing, the most recent is Deep Learning AMI for Amazon Linux 2 is version 26.0, and it’s AMI identifier is [_ami-08e68326c36bf3710_](https://aws.amazon.com/marketplace/search/results?x=0&y=0&searchTerms=ami-08e68326c36bf3710).

![](https://cdn-images-1.medium.com/max/1024/1*AWzrZcLnyTz5vCtiJAezJw.png)

Using this AMI, I fire up a _c5d.4xlarge_ instance. No special settings are required, just make sure you allow SSH access in the Security Group.

![](https://cdn-images-1.medium.com/max/1024/1*94uEjdipHKN2SPasjdDmdQ.png)<figcaption>family, instance name, vCPUs, RAM, storage</figcaption>

Once the instance is up, I ssh to it, and I’m greeted by the familiar Deep Learning AMI banner, tellling me that [Conda](https://docs.conda.io/en/latest/) environments are available for TensorFlow and Apache MXNet.

```
====================================================================
 \_\_| \_\_|\_ )
 \_| ( / Deep Learning AMI (Amazon Linux 2) Version 26.0
\_\_\_|\_\_\_|\_\_\_|
====================================================================

Please use one of the following commands to start the required environment with the framework of your choice:
for MXNet(+Keras2) with Python3 (CUDA 10.1 and Intel MKL-DNN) 
source activate mxnet\_p36
for MXNet(+Keras2) with Python2 (CUDA 10.1 and Intel MKL-DNN) 
source activate mxnet\_p27
**for MXNet(+AWS Neuron) with Python3   
source activate aws\_neuron\_mxnet\_p36**   
for TensorFlow(+Keras2) with Python3 (CUDA 10.0 and Intel MKL-DNN) source activate tensorflow\_p36
for TensorFlow(+Keras2) with Python2 (CUDA 10.0 and Intel MKL-DNN) source activate tensorflow\_p27
**for TensorFlow(+AWS Neuron) with Python3  
source activate aws\_neuron\_tensorflow\_p36**  
for TensorFlow 2(+Keras2) with Python3 (CUDA 10.0 and Intel MKL-DNN) ssource activate tensorflow2\_p36
for TensorFlow 2(+Keras2) with Python2 (CUDA 10.0 and Intel MKL-DNN) ssource activate tensorflow2\_p27
for PyTorch with Python3 (CUDA 10.1 and Intel MKL) 
source activate pytorch\_p36
for PyTorch with Python2 (CUDA 10.1 and Intel MKL)
source activate pytorch\_p27
for Chainer with Python2 (CUDA 10.0 and Intel iDeep)
source activate chainer\_p27
for Chainer with Python3 (CUDA 10.0 and Intel iDeep)
source activate chainer\_p36
for base Python2 (CUDA 10.0)
source activate python2
for base Python3 (CUDA 10.0) 
source activate python3

Official Conda User Guide: [https://docs.conda.io/projects/conda/en/latest/user-guide/](https://docs.conda.io/projects/conda/en/latest/user-guide/)
AWS Deep Learning AMI Homepage: [https://aws.amazon.com/machine-learning/amis/](https://aws.amazon.com/machine-learning/amis/)
Developer Guide and Release Notes: [https://docs.aws.amazon.com/dlami/latest/devguide/what-is-dlami.html](https://docs.aws.amazon.com/dlami/latest/devguide/what-is-dlami.html)
Support: [https://forums.aws.amazon.com/forum.jspa?forumID=263](https://forums.aws.amazon.com/forum.jspa?forumID=263)
For a fully managed experience, check out Amazon SageMaker at [https://aws.amazon.com/sagemaker](https://aws.amazon.com/sagemaker)
**When using INF1 type instances, please update regularly using the instructions at:** [**https://github.com/aws/aws-neuron-sdk/tree/master/release-notes**](https://github.com/aws/aws-neuron-sdk/tree/master/release-notes)
====================================================================
```

I activate the appropriate environment, which provides all required dependencies.

> For the rest of this post, any shell command prefixed by _(aws\_neuron\_tensorflow\_p36)_ should be run inside that Conda environment.

```
$ source activate aws\_neuron\_tensorflow\_p36
(aws\_neuron\_tensorflow\_p36) $
```

Next, I upgrade the _tensorflow-neuron_ package.

```
$ conda install numpy=1.17.2 --yes --quiet
$ conda update tensorflow-neuron
```

We’re now ready to fetch a model and compile it.

#### Compiling a model

The code below fetches a [ResNet50](https://arxiv.org/abs/1512.03385) image classification model pretrained on the [ImageNet](http://image-net.org) dataset, and stores it in the _resnet50_ directory.

Then, it compiles it for Inferentia. I highlighted the single line of code required: everything else is vanilla TensorFlow. Then, the compiled model is saved in the _ws\_resnet50_ directory, and in a ZIP file for easy copy to an _inf1_ instance.

```
import os
import time
import shutil
import tensorflow as tf
**import tensorflow.neuron as tfn**
import tensorflow.compat.v1.keras as keras
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess\_input

# Create a workspace
WORKSPACE = './ws\_resnet50'
os.makedirs(WORKSPACE, exist\_ok=True)

# Prepare export directory (old one removed)
model\_dir = os.path.join(WORKSPACE, 'resnet50')
compiled\_model\_dir = os.path.join(WORKSPACE, 'resnet50\_neuron')
shutil.rmtree(model\_dir, ignore\_errors=True)
shutil.rmtree(compiled\_model\_dir, ignore\_errors=True)

# Instantiate Keras ResNet50 model
keras.backend.set\_learning\_phase(0)
keras.backend.set\_image\_data\_format('channels\_last')

model = ResNet50(weights='imagenet')

# Export SavedModel
tf.saved\_model.simple\_save(
    session = keras.backend.get\_session(),
    export\_dir = model\_dir,
    inputs = {'input': model.inputs[0]},
    outputs = {'output': model.outputs[0]})

**# Compile using Neuron  
tfn.saved\_model.compile(model\_dir, compiled\_model\_dir)**  

# Prepare SavedModel for uploading to Inf1 instance
shutil.make\_archive('./resnet50\_neuron', 'zip', WORKSPACE, 'resnet50\_neuron')
```

That one [API](https://github.com/aws/aws-neuron-sdk/blob/master/docs/tensorflow-neuron/api-compilation-python-api.md) is all it takes! Impressive.

> Power users will enjoy reading about the CLI compiler, [neuron-cc](https://github.com/aws/aws-neuron-sdk/tree/master/docs/neuron-cc).

Running this code produces the expected output.

```
(aws\_neuron\_tensorflow\_p36) $ python compile\_resnet.py

<output removed>
Downloading data from [https://github.com/keras-team/keras-applications/releases/download/resnet/resnet50\_weights\_tf\_dim\_ordering\_tf\_kernels.h5](https://github.com/keras-team/keras-applications/releases/download/resnet/resnet50_weights_tf_dim_ordering_tf_kernels.h5)
102973440/102967424 [==============================] - 2s 0us/step
<output removed>
INFO:tensorflow:fusing subgraph neuron\_op\_d6f098c01c780733 with neuron-cc
INFO:tensorflow:Number of operations in TensorFlow session: 4638
INFO:tensorflow:Number of operations after tf.neuron optimizations: 556
INFO:tensorflow:Number of operations placed on Neuron runtime: 554
INFO:tensorflow:Successfully converted ./ws\_resnet50/resnet50 to ./ws\_resnet50/resnet50\_neuron
```

Then, I simply copy the ZIP file to an Amazon S3 bucket, probably the easiest way to share it with _inf1_ instances used for inference.

```
$ ls \*.zip
resnet50\_neuron.zip
$ aws s3 mb s3://jsimon-inf1-useast1 
$ aws s3 cp resnet50\_neuron.zip s3://jsimon-inf1-useast1
upload: ./resnet50\_neuron.zip to s3://jsimon-inf1-useast1/resnet50\_neuron.zip
```

Alright, let’s fire up one of these babies.

#### Predicting on Inferentia with TensorFlow

Using the same AMI as above, I launch an _inf1.xlarge_ instance.

![](https://cdn-images-1.medium.com/max/1024/1*hMip9-P6o7RaPt46Qk9csg.png)<figcaption>family, instance name, vCPUs, RAM, storage</figcaption>

Once this instance is up, I ssh to it, and I can view some properties using the [_neuron-ls_](https://github.com/aws/aws-neuron-sdk/tree/master/docs/neuron-tools) CLI tool.

![](https://cdn-images-1.medium.com/max/1024/1*nDcXAWoafSbF6Mcc0mi-DQ.png)

4 NeuronCores, as expected. The ‘east’ and ‘west’ columns show connections to other Inferentia chips: as this instance only has one, they’re empty here.

Next, I retrieve the compiled model from my S3 bucket, and extract it. I also retrieve a test image.

```
$ aws s3 cp s3://jsimon-inf1-useast1/resnet50\_neuron.zip .
download: s3://jsimon-inf1-useast1/resnet50\_neuron.zip to resnet50\_neuron.zip
$ unzip resnet50\_neuron.zip
Archive: resnet50\_neuron.zip
 creating: resnet50\_neuron/
 creating: resnet50\_neuron/variables/
 inflating: resnet50\_neuron/saved\_model.pb
$ curl -O [https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten\_small.jpg](https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg)
```

Using the code below, I load and transform the test image. I then load the compiled model, and use it to classify the image.

```
import os
import time
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications import resnet50

tf.keras.backend.set\_image\_data\_format('channels\_last')

# Create input from image
img\_sgl = image.load\_img('kitten\_small.jpg', target\_size=(224, 224))
img\_arr = image.img\_to\_array(img\_sgl)
img\_arr2 = np.expand\_dims(img\_arr, axis=0)
img\_arr3 = resnet50.preprocess\_input(img\_arr2)

# Load model
COMPILED\_MODEL\_DIR = './resnet50\_neuron/'
predictor\_inferentia = tf.contrib.predictor.from\_saved\_model(COMPILED\_MODEL\_DIR)

# Run inference
model\_feed\_dict={'input': img\_arr3}
infa\_rslts = predictor\_inferentia(model\_feed\_dict);

# Display results
print(resnet50.decode\_predictions(infa\_rslts["output"], top=5)[0])
```

Can you guess how many lines of Inferentia specific code are present here? The answer is **zero**. We seamlessly use the [_tf.contrib.predictor_](https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/contrib/predictor) API. Woohoo!

Running this code produces the expected output, and we see the top 5 classes for the image.

```
(aws\_neuron\_tensorflow\_p36) $ python infer\_resnet50.py

<output removed>
[('n02123045', 'tabby', 0.6918919), ('n02127052', 'lynx', 0.12770271), ('n02123159', 'tiger\_cat', 0.08277027), ('n02124075', 'Egyptian\_cat', 0.06418919), ('n02128757', 'snow\_leopard', 0.009290541)]
```

Now let’s see how we can deploy a compiled model using [TensorFlow Serving,](https://www.tensorflow.org/tfx/guide/serving) which is a very good option for production deployments.

#### Predicting on Inferentia with TensorFlow Serving

First, we need to package the model properly, and move it to a directory reflecting it’s version. We have only one here, so let’s move the saved model to a directory named ‘1’.

```
$ pwd
/home/ec2-user/resnet50\_neuron
$ mkdir 1
$ mv \* 1
```

Now, we can launch TensorFlow Serving, and load the compiled model. Once again, this is vanilla TensorFlow.

```
(aws\_neuron\_tensorflow\_p36) $ tensorflow\_model\_server\_neuron 
--model\_name=resnet50 
--model\_base\_path=/home/ec2-user/resnet50\_neuron 
--port=8500

2019–12–13 16:16:27.704882: I tensorflow\_serving/core/loader\_harness.cc:87] **Successfully loaded servable version {name: resnet50 version: 1}**
2019–12–13 16:16:27.706241: I tensorflow\_serving/model\_servers/server.cc:353] Running gRPC ModelServer at 0.0.0.0:8500
```

Once TensorFlow Serving is up and running, we can use the script below to load a test image, and send it for prediction. At the risk of repeating myself… this is vanilla TensorFlow :)

```
import numpy as np
import grpc
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess\_input
from tensorflow.keras.applications.resnet50 import decode\_predictions
from tensorflow\_serving.apis import predict\_pb2
from tensorflow\_serving.apis import prediction\_service\_pb2\_grpc

if \_\_name\_\_ == '\_\_main\_\_':
    chan = grpc.insecure\_channel('localhost:8500')
    stub = prediction\_service\_pb2\_grpc.PredictionServiceStub(chan)

    img\_file="kitten\_small.jpg",      
    img = image.load\_img(img\_file, target\_size=(224, 224))
    img\_array = preprocess\_input(image.img\_to\_array(img)[None, ...])

    request = predict\_pb2.PredictRequest()
    request.model\_spec.name = 'resnet50\_inf1'
    request.inputs['input'].CopyFrom(
        tf.contrib.util.make\_tensor\_proto(
            img\_array, shape=img\_array.shape)
    )

    result = stub.Predict(request)
    prediction = tf.make\_ndarray(result.outputs['output'])
    print(decode\_predictions(prediction))
```

Running this code produces the expected output, and we see the top 5 classes for the image.

```
(aws\_neuron\_tensorflow\_p36) $ python tfserving\_resnet50.py

<output removed>
[[(‘n02123045’, ‘tabby’, 0.6918919), (‘n02127052’, ‘lynx’, 0.12770271), (‘n02123159’, ‘tiger\_cat’, 0.08277027), (‘n02124075’, ‘Egyptian\_cat’, 0.06418919), (‘n02128757’, ‘snow\_leopard’, 0.009290541)]]
```

#### Diving deeper

That’s it for today. I hope I gave you a clear introduction to AWS Inferentia, and how easy it is to use it! All it took is **one line of code** to compile our model.

If you’d like to dive deeper, I highly recommend the excellent [workshop](https://github.com/awshlabs/reinvent19Inf1Lab) delivered at re:Invent by my colleague Wenming Ye. One of the labs shows you how to compile a 32-bit floating point (FP32) ResNet50 model to 16-bit floating point (FP16). By reducing arithmetic complexity, this technique is known to improve performance while preserving accuracy. Indeed, on an _inf1.2xlarge_ instance, the FP16 model delivers an impressive **1,500 image classifications per second**!

As always, thank you for reading. Happy to answer questions here or on [Twitter](https://twitter.com/julsimon).

Inferentia NeuronCores are totally hardcore. They destroy everything \m/

{% youtube DBwgX8yBqsw %}