# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: conda_tensorflow_p36
#     language: python
#     name: conda_tensorflow_p36
# ---

# ## TensorFlow BYOM: 独自の学習スクリプトを用いたモデルの学習と、デプロイ手段の比較
#
# 本ハンズオンでは [TensorFlow MNIST distributed training notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/tensorflow_distributed_mnist/tensorflow_distributed_mnist.ipynb) と同様に [MNIST dataset](http://yann.lecun.com/exdb/mnist/) を TensorFlow で分散学習を行います。その後、CPU () を用いた推論インスタンスへのデプロイ、Elastic Inference を用いた推論インスタンスへのデプロイ、SageMaker Neo を用いてモデルをコンパイルした上での推論インスタンスへのデプロイを行います。
#
# ---
# ## コンテンツ
#
# 1. [環境のセットアップ](#1.環境のセットアップ)
# 1. [学習データの準備](#2.学習データの準備)
# 1. [分散学習用のスクリプトの準備](#3.分散学習での学習スクリプトの準備)
# 1. [TensorFlow Estimator を利用して学習ジョブを作成する](#4.TensorFlowEstimatorを利用して学習ジョブを作成する)
# 1. [CPUインスタンスを用いた推論](#5.CPUインスタンスを用いた推論)
# 1. [Elastic Inference を使った推論](#6.ElasticInferenceを使った推論)
# 1. [SageMaker Neo をでコンパイルしたモデルによる推論](#7.SageMakerNeoをでコンパイルしたモデルによる推論)
# 1. [エンドポイントの削除](#8.エンドポイントの削除)
# ---

# ## 1. 環境のセットアップ
#
# まずは環境のセットアップを行いましょう。

# +
import os
import io
import time
import sagemaker
from sagemaker import get_execution_role

import numpy as np
from tensorflow.examples.tutorials.mnist import input_data

sagemaker_session = sagemaker.Session()
role = get_execution_role()
# -

# # 2.学習データの準備
# MNIST データセットのダウンロードし。学習、評価、テスト用のそれぞれのデータへ分割します。

# +
import utils
from tensorflow.contrib.learn.python.learn.datasets import mnist
import tensorflow as tf

data_sets = mnist.read_data_sets('data', dtype=tf.uint8, reshape=False, validation_size=5000)

utils.convert_to(data_sets.train, 'train', 'data')
utils.convert_to(data_sets.validation, 'validation', 'data')
utils.convert_to(data_sets.test, 'test', 'data')

mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
# -

# データを Amazon S3 上へアップロードします。

inputs = sagemaker_session.upload_data(path='data', key_prefix='data/DEMO-mnist')

# # 3.分散学習での学習スクリプトの準備
#
# このチュートリアルのトレーニングスクリプトは、TensorFlowの公式の[CNN MNISTの例](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/layers/cnn_mnist.py) をベースに作成されました。 SageMaker から渡された `` model_dir`` パラメーターを処理するように変更しています。 これは、分散学習時のデータ共有、チェックポイント、モデルの永続保存などに使用できるS3パスです。 また、トレーニング関連の変数を扱うために、引数をパースする関数も追加しました。
#
# スクリプト全体は次のとおりです。

# !cat 'mnist.py'

# # 4.TensofFlow Estimator による学習ジョブの作成
#
# `sagemaker.tensorflow.TensorFlow`　estimator は、TensorFlow コンテナの指定、学習・推論スクリプトの S3 へのアップロード、および SageMaker トレーニングジョブの作成を行います。ここでいくつかの重要なパラメーターを呼び出しましょう。
#
# `distributions` は、分散トレーニング設定を構成するために使用されます。インスタンスのクラスターまたは複数の GPU をまたいで分散学習を行う場合にのみ必要です。ここでは、分散トレーニングスキーマとしてパラメーターサーバーを使用しています。 SageMaker トレーニングジョブは同種のクラスターで実行されます。 SageMaker セットアップでパラメーターサーバーのパフォーマンスを向上させるために、クラスター内のすべてのインスタンスでパラメーターサーバーを実行するため、起動するパラメーターサーバーの数を指定する必要はありません。スクリプトモードは、[Horovod](https://github.com/horovod/horovod) による分散トレーニングもサポートしています。 `distributions` の設定方法に関する詳細なドキュメントは[こちら](https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/tensorflow#distributed-training) をご参照ください。

# +
from sagemaker.tensorflow import TensorFlow

mnist_estimator = TensorFlow(entry_point='mnist.py',
                             role=role,
                             framework_version='1.11.0',
                             training_steps=1000, 
                             evaluation_steps=100,
                             train_instance_count=2,
                             train_instance_type='ml.p2.xlarge')

mnist_estimator.fit(inputs)
# -

# #### ``fit`` による学習ジョブの実行
#
# 学習ジョブを開始するには、`estimator.fit（training_data_uri）` を呼び出します。
#
# ここでは、S3 ロケーションが入力として使用されます。 `fit` は、`training` という名前のデフォルトチャネルを作成します。これは、このS3ロケーションを指します。トレーニングスクリプトでは、 `SM_CHANNEL_TRAINING` に保存されている場所からトレーニングデータにアクセスできます。 `fit`は、他のいくつかのタイプの入力も受け入れます。詳細については、APIドキュメント[こちら](https://sagemaker.readthedocs.io/en/stable/estimators.html#sagemaker.estimator.EstimatorBase.fit) を参照してください。
#
# トレーニングが開始されると、TensorFlow コンテナは mnist.py を実行し、スクリプトの引数として　estimator から`hyperparameters` と `model_dir` を渡します。この例では、estimator 内で定義していないハイパーパラメーターは渡されず、 `model_dir` のデフォルトは `s3://<DEFAULT_BUCKET>/<TRAINING_JOB_NAME>` であるため、スクリプトの実行は次のようになります。
# ```bash
# python mnist.py --model_dir s3://<DEFAULT_BUCKET>/<TRAINING_JOB_NAME>
# ```
# トレーニングが完了すると、トレーニングジョブは保存されたモデルを TensorFlow serving にアップロードします。

# # 5.CPUインスタンスを用いた推論

# #### 5.1 m4.xlarge CPU インスタンスのデプロイ
# `deploy（）`メソッドは SageMaker モデルを作成します。このモデルはエンドポイントにデプロイされ、リアルタイムで予測リクエストを処理します。まずは`ml.m4.xlarge` インスタンスを用いた推論インスタンスをデプロイします。

predictor_cpu = mnist_estimator.deploy(initial_instance_count=1,
                                     instance_type='ml.m4.xlarge',
                                     wait=False)

# #### 5.2 m4.xlarge CPU インスタンスでのリアルタイム推論
# モデルをデプロイしたエンドポイントに対して、データを入力し、推論結果を得ます。

# +
process_times = []
for i in range(20):
    data = mnist.test.images[i].tolist()
    
    start = time.time()
    predict_response = predictor_cpu.predict(data)
    process_time = time.time() - start
    process_times.append(process_time)
    
    label = np.argmax(mnist.test.labels[i])
    prediction = predict_response['outputs']['classes']['int64_val'][0]
    print("label is {}, prediction is {}".format(label, prediction))
    
print("推論実行時間の平均は {} です。".format(np.array(process_times).mean()))
# -

# # 6. Elastic Inference を使った推論

# #### 6.1 Elastic Inference を用いた推論インスタンスへのデプロイ
# Amazon Elastic Inference (EI) を使用することで、Amazon SageMaker がホストするモデルとしてデプロイされている深層学習モデルからのスループットを高速化し、リアルタイムの推論を得るためのレイテンシーを短縮することができます。さらに、エンドポイントに GPU インスタンスを使用するコストと比較して、大幅なコストダウンを実現できます。`accelerator_type` を指定することで Elastic Inference を使った推論インスタンスをデプロイできます。

# +
from sagemaker.tensorflow.serving import Model

saved_model = mnist_estimator.model_data

# モデルの作成
estimator_ei = Model(model_data=saved_model,
                     role=role,
                     framework_version='1.11')

# モデルのデプロイ
predictor_ei = estimator_ei.deploy(initial_instance_count=1,
                                   instance_type='ml.m4.xlarge',
                                   accelerator_type='ml.eia1.xlarge',
                                   wait=False)
# -

# #### 6.2 Elastic Inference を使ったリアルタイム推論
# 同様に、推論結果を得ます。

# +
process_times = []

for i in range(20):
    data = mnist.test.images[i].tolist()
    
    start = time.time()
    predict_response = predictor_ei.predict(data)
    process_time = time.time() - start
    process_times.append(process_time)
    
    label = np.argmax(mnist.test.labels[i])
    prediction = predict_response['predictions'][0]['classes']
    print("label is {}, prediction is {}".format(label, prediction))
    
print("推論実行時間の平均は {} です。".format(np.array(process_times).mean()))
# -

# # 7. SageMaker Neo をでコンパイルしたモデルによる推論

# #### 7.1 SageMaker Neo でのモデルのコンパイル
#
# SageMaker Neo API を用いれば学習済みのモデルを特定のハードウェア用に最適化することが可能です。`compile_model()` 関数を呼ぶ際に、ターゲットとなるインスタンスファミリー (ここでは M4) とコンパイル済みのモデルを保存する S3 バケットを指定します。
#
# ** [重要] もし以下のコマンドが permission error になる場合、ノートブックの上部にスクロールして `get_execution_role()` により返される execution role の値を確認して下さい。このロールには ``output_path`` で指定される S3 バケットへのアクセス権限が必要です。

output_path = '/'.join(mnist_estimator.output_path.split('/')[:-1])
optimized_estimator = mnist_estimator.compile_model(target_instance_family='ml_m4', 
                              input_shape={'data':[1, 784]},  # Batch size 1, 3 channels, 224x224 Images.
                              output_path=output_path,
                              framework='tensorflow', framework_version='1.11.0')

# #### 7.2 SageMAker Neo でコンパイルしたモデルのデプロイ 
#
# SageMaker Neo によってコンパイルされたモデルを推論インスタンスへデプロイします。

# +
optimized_predictor = optimized_estimator.deploy(initial_instance_count = 1,
                                                 instance_type = 'ml.m4.xlarge',
                                                 wait=False)

def numpy_bytes_serializer(data):
    f = io.BytesIO()
    np.save(f, data)
    f.seek(0)
    return f.read()

optimized_predictor.content_type = 'application/vnd+python.numpy+binary'
optimized_predictor.serializer = numpy_bytes_serializer
# -

# #### 7.3 SageMaker Neo でコンパイルしたモデルでのリアルタイム推論

# +
for i in range(20):
    data = mnist.test.images[i]
    
    # Invoke endpoint with image
    start = time.time()
    predict_response = optimized_predictor.predict(data)
    process_time = time.time() - start
    process_times.append(process_time)
        
    label = np.argmax(mnist.test.labels[i])
    prediction = np.array(predict_response).argmax()
    print("label is {}, prediction is {}".format(label, prediction))
    
print("推論実行時間の平均は {} です。".format(np.array(process_times).mean()))
# -

# # 8. エンドポイントの削除
#
# このノートブックによって作られたリソースを削除していい場合、以下のセルを実行してください。このコマンドは上で作成したエンドポイントを削除して意図しない請求を防ぐことができます。
# (必要であれば、このノートブック自体を走らせているノートブックインスタンスも SageMaker のマネージメントコンソールから停止させて下さい。)

sagemaker.Session().delete_endpoint(predictor_cpu.endpoint)
sagemaker.Session().delete_endpoint(predictor_ei.endpoint)
sagemaker.Session().delete_endpoint(optimized_predictor.endpoint)



StarcoderdataJupyterScriptsDedupFiltered