TensorFlow-Neuron と AWS Neuron Compiler の使用

このチュートリアルでは、 AWS Neuron コンパイラを使用して Keras ResNet-50 モデルをコンパイルし、保存されたモデルとして SavedModel 形式でエクスポートする方法を示します。この形式は、一般的な TensorFlow モデル交換可能な形式です。また、入力例を使用して Inf1 インスタンスで推論を実行する方法も学習します。 

Neuron SDK の詳細については、AWS Neuron SDK のドキュメントを参照してください。

前提条件

このチュートリアルを使用する前に、AWS Neuron を使用した DLAMI インスタンスの起動 の設定ステップを完了しておく必要があります。また、深層学習および DLAMI の使用にも精通している必要があります。

Conda 環境のアクティブ化

次のコマンドを使用して TensorFlow-Neuron conda 環境をアクティブ化します。

source activate aws_neuron_tensorflow_p36

現在の Conda 環境を終了するには、次のコマンドを実行します。

source deactivate

Resnet50 コンパイル

次の内容を含む tensorflow_compile_resnet50.py という Python スクリプトを作成します。この Python スクリプトは Keras ResNet50 モデルをコンパイルし、保存されたモデルとしてエクスポートします。

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)

# Prepare SavedModel for uploading to Inf1 instance
shutil.make_archive(compiled_model_dir, 'zip', WORKSPACE, 'resnet50_neuron')

次のコマンドを使用してモデルをコンパイルします。

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

ResNet50 推論

次の内容を含む tensorflow_infer_resnet50.py という Python スクリプトを作成します。このスクリプトは、以前にコンパイルされた推論モデルを使用して、ダウンロードしたモデルに対して推論を実行します。

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 = './ws_resnet50/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 TensorFlow Serving の使用