在2019年底进行的年度AWS re:Invent大会上，最令人激动的新发布之一就是正式推出AWS 自己设计和打造的Inferentia芯片，以更低成本提供更高性能的机器学习推理能力。
机器学习推理是使用经过训练的模型做出预测的流程。经统计，在机器学习应用程序中，推理最多能够占到总成本的 90%。其原因有二：首先，独立 GPU 实例通常专为模型训练而设计，而非用于推理。虽然训练作业可并行批量处理大量数据样本，但推理作业往往会实时处理单个输入，因而仅占用少量 GPU 计算。这使得独立 GPU 推理成本高且效率低。另一方面，独立的CPU实例并非专门用于矩阵运算，因此对于深度学习推理而言通常太慢。 其次，不同的模型具有不同的CPU，GPU和内存要求。 对一种资源进行优化可能导致对其他资源的利用不足和更高的成本。
对于机器学习推理而言，AWS Inferentia 芯片的优势可以总结为：高性能，低延迟，高易用。每个 AWS Inferentia 芯片有4个处理核心Neuron Core，能在低功率下支持高达 128 TOPS的性能。AWS Inferentia 能够支持 FP16、BF16 和 INT8 数据类型，并且可以接收FP32的训练模型输入，并使用 BF16高速运行该模型。AWS Inferentia 具有大容量的片上存储，可用于缓存大型模型，从而无需将它们存储到片外，使得Neuron Core可以对模型进行高速访问，对于降低推理延迟具有显著影响。AWS Inferentia附带了AWS Neuron SDK，可以使用AWS Inferentia在流行框架中创建和训练复杂的神经网络模型。Neuron由编译器、运行时和分析工具组成，并预先集成到流行的机器学习框架中，包括TensorFlow、Pytorch和MXNet。
基于AWS Inferentia芯片，AWS在云端提供了EC2 Inf1实例，将 Inferentia 芯片与AWS Nitro虚拟化管理程序，最新的第二代定制 Intel Xeon 可扩展处理器，和高达 100Gbps 的网络相结合，以实现高吞吐量推理。这一配置使 Inf1 实例提供 Amazon EC2 G4 实例 3 倍的吞吐量和40%推理成本的降低(基于使用TensorFlow端到端运行BERT模型的结果)。Inf1实例的规格如下表所示。
| 实例规格 |
Inferentia芯片 |
vCPU |
内存 |
EBS 带宽 |
网络带宽 |
| inf1.xlarge |
1 |
4 |
8 GiB |
最高4.75 Gbps |
最高25 Gbps |
| inf1.2xlarge |
1 |
8 |
16 GiB |
最高4.75 Gbps |
最高25 Gbps |
| inf1.6xlarge |
4 |
24 |
48 GiB |
4.75 Gbps |
25 Gbps |
| inf1.24xlarge |
16 |
96 |
192 GiB |
19 Gbps |
100 Gbps |
使用AWS Inferentia进行机器学习的典型流程如下图所示。通常，我们会在GPU实例上训练机器学习模型，之后我们把训练好的模型利用 AWS Neuron进行重新编译，然后将此模型交付给Inf1实例 (AWS Inferentia芯片)执行机器学习推理。
下面，我们通过一个样例来看如何使用AWS Neuron在AWS Inferentia芯片上进行机器学习推理。在这里，我们会用到AWS Deep Learning AMI，它可以在Amazon EC2上一键式安装机器学习的各种框架工具，从而加快在云中进行机器学习的速度。最新的Deep Learning AMI已支持AWS Inferentia，并附带了AWS Neuron开发所需的工具包 (Neuron SDK)。
我们在us-west-2 区域以Deep Learning AMI (Ubuntu 18.04) Version 27.0 – ami-008d8ed4bd7dc2485做为系统镜像，启动Inf1.2xlarge实例。在启动Inf1实例后，我们使用以下命令升级到最新版的Neuron运行时环境和工具包。
$ sudo apt-get update
$ sudo apt-get install aws-neuron-runtime
$ sudo apt-get install aws-neuron-tools
此时，我们就能够用neuron-ls命令显示Inf1实例上的Inferentia芯片的数量和信息。
$ neuron-ls
+--------------+---------+--------+-----------+-----------+------+------+
| PCI BDF | LOGICAL | NEURON | MEMORY | MEMORY | EAST | WEST |
| | ID | CORES | CHANNEL 0 | CHANNEL 1 | | |
+--------------+---------+--------+-----------+-----------+------+------+
| 0000:00:1f.0 | 0 | 4 | 4096 MB | 4096 MB | 0 | 0 |
+--------------+---------+--------+-----------+-----------+------+------+
在输出的表格中，第一列显示PCI总线设备功能ID。第二列显示分配给设备的逻辑ID。这个逻辑ID在Neuron运行时守护进程 (Runtime Daemon)（Neuron rtd）配置期间使用。第三列显示可用的Neuron Core数量。最后两列显示与任何其他Inferentia设备的连接。因为我们选择的Inf1.2xlarge只有一个Inferentia芯片，所以这两列为空 （如果我们启动Inf1.6xlarge或者Inf1.24xlarge这两种带有多个Inferentia芯片的实例，这两列会显示Inferentia芯片的互联信息）。
在搭建好环境之后，我们来看一下如何使用AWS Neuron和Tensorflow进行深度学习的推理。整个过程会分为两个阶段，编译阶段和推理阶段。
在第一阶段，我们将编译TensorFlow Neuron中预训练的Keras ResNet50模型，并将其导出为SavedModel，一种TensorFlow模型的交换格式。
首先，我们来激活TensorFlow-Neuron conda虚拟环境 。使用虚拟环境能够确保在使用深度学习框架时对包的管理有最佳的灵活性，这也是深度学习的一个最佳实践。
$ source activate aws_neuron_tensorflow_p36
之后我们使用以下命令更新Neuron包到最新版本
$ conda update tensorflow-neuron
接下来，我们创建一个名为tensorflow_compile_resnet50.py的Python脚本，此Python脚本使用AWS Neuron编译Keras ResNet50模型并将其导出为TensorFlow SavedModel格式。
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)
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, compiler_args =['--num-neuroncores', '4'])
# Prepare SavedModel for uploading to Inf1 instance
shutil.make_archive(compiled_model_dir, 'zip', WORKSPACE, 'resnet50_neuron')
这里有两点需要注意。第一，Resnet50的模型有2000多万个参数。使用FP32格式，每个参数4字节。AWS Neuron编译器自动将它们转换为BF16，每个参数2字节，这是一种更有效的数据格式，Neuron Core在硬件上原生支持这种格式；第二，由于我们启用的inf1.2xlarge实例有4个Neuron Core, 设置compiler_args =[‘–num-neuroncores’, ‘4’]，会将模型定位并优化于4个Neuron Core上运行。此参数的默认值为1，我们需要确保参数设置为启动的Inf1实例上的Neuron Core的数量。
随后，我们使用以下命令编译模型：
$ python tensorflow_compile_resnet50.py
编译过程大约需要几分钟时间，输出应包含如下信息：
...
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
...
编译后，保存的模型以压缩格式存储在在ws_resnet50/resnet50_neuron.zip。我们使用以下命令解压缩模型
$ unzip ws_resnet50/resnet50_neuron.zip -d .
在第二阶段，我们会使用在第一阶段编译的推理模型在示例图像上运行推理。
首先，我们下载一幅用于推理的示例图像
$ curl -O https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg
之后，我们创建一个名为tensorflow_infer_resnet50.py的Python脚本，该脚本包含以下内容，使用在第一阶段编译好的模型在AWS Inferentia上进行推理。
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications import resnet50
# 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])
这样，我们就可以使用以下命令运行推理操作：
$ python tensorflow_infer_resnet50.py
输出应该如下所示：
...
[('n02123045', 'tabby', 0.6918919),
('n02127052', 'lynx', 0.12770271),
('n02123159', 'tiger_cat', 0.08277027),
('n02124075', 'Egyptian_cat', 0.06418919),
('n02128757', 'snow_leopard', 0.009290541)]
小结
通过上述的样例，我们看到了使用AWS Neuron SDK在AWS Inferentia芯片上进行机器学习推理的全过程，也希望借此能够帮助大家把更多的机器学习应用在云端部署以获得最佳的体验。
本篇作者