亚马逊云科技为机器学习推理开发的专用芯片，你不想体验一下？

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

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

2019年底进行的年度亚马逊云科技 (Amazon Web Services）re:Invent 大会上，最令人激动的新发布之一就是正式推出亚马逊云科技自己设计和打造的 Inferentia 芯片，以更低成本提供更高性能的机器学习推理能力。

机器学习推理是使用经过训练的模型做出预测的流程。经统计，在机器学习应用程序中，推理最多能够占到总成本的90%。其原因有二：

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

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

• 每个 Amazon Inferentia 芯片有4个处理核心 Neuron Core，能在低功率下支持高达 128 TOPS 的性能。
• Amazon Inferentia 支持 FP16、BF16 和 INT8 数据类型，并且可接收 FP32 的训练模型输入，并使用 BF16 高速运行该模型。
• Amazon Inferentia 具有大容量片上存储，可用于缓存大型模型，从而无需将它们存储到片外，这使得 Neuron Core 可对模型进行高速访问，对于降低推理延迟具有显著影响。
• Amazon Inferentia 附带了 Amazon Neuron SDK，可使用 Amazon Inferentia 在流行框架中创建和训练复杂的神经网络模型。
• Neuron 由编译器、运行时和分析工具组成，并预先集成到流行的机器学习框架中，包括 TensorFlow、Pytorch 和 MXNet。

基于 Amazon Inferentia 芯片，Amazon 在云端提供了 EC2 Inf1 实例，将 Inferentia 芯片与 Amazon Nitro 虚拟化管理程序、最新的第二代定制 Intel Xeon 可扩展处理器，以及高达 100Gbps 的网络相结合，以实现高吞吐量推理。这一配置使 Inf1 实例提供了 Amazon EC2 G4 实例3倍的吞吐量和40%推理成本的降低（基于使用 TensorFlow 端到端运行 BERT 模型的结果)。

Inf1 实例的规格如下表所示。

使用 Amazon Inferentia 进行机器学习的典型流程如下图所示。通常，我们会在 GPU 实例上训练机器学习模型，之后把训练好的模型利用 Amazon Neuron 重新编译，然后将此模型交付给 Inf1 实例（Amazon Inferentia 芯片）执行机器学习推理。

下面，我们通过一个样例来看看如何使用 Amazon Neuron 在 Amazon Inferentia 芯片上进行机器学习推理。

在这里，我们会用到 Amazon Deep Learning AMI，它可以在 Amazon EC2 上一键式安装机器学习的各种框架工具，从而加快在云中进行机器学习的速度。最新的 Deep Learning AMI 已支持 Amazon Inferentia，并附带 Amazon 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 芯片的互联信息）。

搭建好环境后，来看一下如何使用 Amazon 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 脚本使用 Amazon 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字节。Amazon 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 脚本，该脚本包含以下内容，使用在第一阶段编译好的模型在 Amazon 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)]

 

总结

通过上述样例，我们看到了使用 Amazon Neuron SDK 在 Amazon Inferentia 芯片上进行机器学习推理的全过程，也希望借此能够帮助大家把更多的机器学习应用在云端部署以获得最佳的体验。