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 芯片上進行機器學習推理的全過程，也但願藉此可以幫助你們把更多的機器學習應用在雲端部署以得到最佳的體驗。