コンパイルとInferentia上での実行

章4.1でSavedModel形式でエクスポートした事前学習済みの Keras ResNet-50 モデルを Inferentia推論チップ向けにコンパイルし、Neuronコア上での画像分類を実行します。 コンパイル後のモデルの出力形式もSavedModelです。

SavedModel は、モデルの重みとネットワーク構造が一緒に保存された TensorFlow モデルと交換可能な形式です。　

Step 1. コンパイル用 Python スクリプトを作成

以下の内容で compile_resnet50.py というファイル名のコンパイル用 Python スクリプトを作成します。

import shutil
import tensorflow.neuron as tfn

model_dir = 'resnet50'

# Prepare export directory (old one removed)
compiled_model_dir = 'resnet50_neuron'
shutil.rmtree(compiled_model_dir, ignore_errors=True)

# Compile using Neuron
tfn.saved_model.compile(model_dir, compiled_model_dir)

Step 2. コンパイル用スクリプトを実行

コンパイルスクリプト を実行します。inf1.2xlarge では~2分程度かかります。

time python compile_resnet50.py  
....
Compiler status PASS
INFO:tensorflow:Number of operations in TensorFlow session: 4638
INFO:tensorflow:Number of operations after tf.neuron optimizations: 876
INFO:tensorflow:Number of operations placed on Neuron runtime: 874
INFO:tensorflow:Successfully converted resnet50 to resnet50_neuron

real    1m9.173s
user    0m56.236s
sys     0m2.773s

Neuron コンパイラでは、事前 (Ahead-Of-Time) コンパイルを行います。実行時 (JIT) コンパイラ や初回推論コンパイルシステムと比較すると、複数のインスタンス上にデプロイする場合において時間を節約できます。Neuron コンパイラは、自動的にオペレータを融合し、 スケジューリングとメモリ管理を行います。

Step 3. 推論実行 Python スクリプトを作成

以下の内容でinfer_resnet50_neuron.py というファイル名の推論実行 Python スクリプトを作成します。CPU上で推論を実行する場合と異なり、Step 2でコンパイルしたモデルを利用します。

import os
import time
import shutil
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input, decode_predictions

# Load model
compiled_model_dir = 'resnet50_neuron'
predictor_inferentia = tf.contrib.predictor.from_saved_model(compiled_model_dir)

# 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 = preprocess_input(img_arr2)

# Run inference, Display results
model_feed_dict={'input': img_arr3}
infa_rslts = predictor_inferentia(model_feed_dict)
print(decode_predictions(infa_rslts["output"], top=5)[0])

Step 4. 推論スクリプトを実行

Neuronコア上でのモデル推論を実行するために、推論実行スクリプトinfer_resnet50_neuron.py を実行します。

python infer_resnet50_neuron.py

CPU上で実行した場合と同様の次の結果が取得されます。

[('n02123045', 'tabby', 0.684492), ('n02127052', 'lynx', 0.1263369), ('n02123159', 'tiger_cat', 0.086898394), ('n02124075', 'Egyptian_cat', 0.067847595), ('n02128757', 'snow_leopard', 0.00977607)]