AWS 为机器学习推理开发的专用芯片，你不想体验一下？

生活中不少经验告诉咱们，不少状况下，一些「专用」的东西，每每比「通用」的在执行具体任务时效果更好，例如为了分担 CPU 负载，实现更强图形处理性能，电脑开始配置专用的图形处理器 GPU。这两年，随着机器学习技术的火热，人们发现，用 GPU 来进行 ML 模型的训练也是一种比较好的作法，速度更快，性能更强。

然而毕竟 GPU 诞生的本意就是用于处理计算机图像的，并不是「AI」专用，那么有没有专供机器学习使用的芯片，能够实现比 GPU 训练更快的速度，以及更低的成本？

2019年末进行的年度 AWS re:Invent 大会上，最使人激动的新发布之一就是正式推出 AWS 本身设计和打造的 Inferentia 芯片，以更低成本提供更高性能的机器学习推理能力。

机器学习推理是使用通过训练的模型作出预测的流程。经统计，在机器学习应用程序中，推理最多可以占到总成本的90%。其缘由有二：

1. 独立 GPU 实例一般专为模型训练而设计，而非用于推理。虽然训练做业可并行批量处理大量数据样本，但推理做业每每会实时处理单个输入，于是仅占用少许 GPU 计算。这使得独立 GPU 推理成本高且效率低。
2. 独立 CPU 实例并不是专门用于矩阵运算，所以对于深度学习推理而言一般太慢。其次，不一样模型对 CPU、GPU 和内存有不一样要求，对一种资源进行优化可能致使对其余资源的利用不足和更高的成本。

对于机器学习推理而言，AWS Inferentia 芯片的优点能够总结为：高性能，低延迟，高易用。

• 每一个 AWS Inferentia 芯片有4个处理核心 Neuron Core，能在低功率下支持高达 128 TOPS 的性能。
• AWS Inferentia 支持 FP1六、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 实例的规格以下表所示。

使用 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 芯片上进行机器学习推理的全过程，也但愿借此可以帮助你们把更多的机器学习应用在云端部署以得到最佳的体验。