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

# # TensorFlow2 と SMDataParallel を使った MNIST のデータ並列分散学習
#
# SMDataParallelは、Amazon SageMaker の新しい機能で、深層学習モデルをより速く、より安価に学習することができます。SMDataParallelは、TensorFlow2、PyTorch、MXNet 用の分散データ並列学習フレームワークです。
#
# このノートブックでは、MNIST データセットを使用して SageMaker で TensorFlow2 と SMDataParallel を使用する方法をご紹介します。
#
# 詳細に関しては以下のリンクをご参照ください。
# 1. [TensorFlow in SageMaker](https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html)
# 2. [SMDataParallelTensorFlow API Specification](https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-use-api.html)
# 3. [Getting started with SMDataParallel on SageMaker] < LINK TO BE ADDED >
#
# **NOTE:** このノートブックでは SageMaker Python SDK v2.X. を使用しています。
#
# ### Dataset
# この例は MNIST データセットを使用しています。MNIST は、手書き数字の分類に広く使われているデータセットです。MNIST データセットは、ラベル付けされた 28x28 ピクセルの手書き数字のグレースケール画像 70,000枚で構成されています。データセットは 60,000枚の学習画像と 10,000枚のテスト画像に分割されています。0 から 9 までの 10個のクラスがあります。
#
# ### SageMaker execution roles
# IAMロール ARN は、トレーニングやデータへのアクセス権を与えるために使用されます。これらの作成方法については、[Amazon SageMaker ロール](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) を参照してください。ノートブックインスタンス、トレーニング、ホスティングに複数のロールを使用する場合は、sagemaker.get_execution_role() を適切な IAM ロールの ARN 文字列に置き換えてください。

pip install sagemaker --upgrade

# +
import sagemaker

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

# ## SMDataParallel を使ったモデルの学習
#
# ### 学習スクリプト
#
# `train_tensorflow_smdataparallel_mnist.py` スクリプトは、SMDataParallel の `DistributedGradientTape` を使用して SageMaker モデルをトレーニングするために必要なコードを提供します。この学習スクリプトは、SageMaker の外部で実行する TensorFlow2 学習スクリプトに非常に似ていますが、SMDataParallel で実行するように修正されています。SMDataParallel の TensorFlow クライアントは、ネイティブの `DistributedGradientTape` に代わるものを提供します。ネイティブTF2 スクリプトでの SMDataParallel の使用方法の詳細については、[SMDataParallel チ ュートリアル](https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-modify-sdp.html#data-parallel-modify-sdp-tf2) を参照してください。

# !pygmentize code/train_tensorflow_smdataparallel_mnist.py

# ### SageMaker TensorFlow Estimator function options
#
# 次のコードブロックでは、異なるインスタンスタイプ、インスタンス数、および分散学習の設定を使用するように、estimator を設定します。また、前のセルで確認した学習スクリプトも指定します。
#
#
# **インスタンスタイプ**
#
# SMDataParallel は以下のインスタンスタイプのみで SageMaker のモデル学習に対応しています。
# 1. ml.p3.16xlarge
# 1. ml.p3dn.24xlarge [Recommended]
# 1. ml.p4d.24xlarge [Recommended]
#
# **インスタンス数**
#
# SMDataParallel によって最高のパフォーマンスを得るには少なくとも 2つのインスタンスを使う必要がありますが、動作確認のために 1つのインスタンスのみを使用することも可能です。
#
# **分散学習の設定**
#
# DDP モードを使用するには `smdistributed dataparallel` を使用するように `distribution` strategy を設定する必要があります。

from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
                        base_job_name='tensorflow2-smdataparallel-mnist',
                        source_dir='code',
                        entry_point='train_tensorflow_smdataparallel_mnist.py',
                        role=role,
                        py_version='py37',
                        framework_version='2.3.1',
                        # For training with multinode distributed training, set this count. Example: 2
                        instance_count=2,
                        # For training with p3dn instance use - ml.p3dn.24xlarge
                        instance_type= 'ml.p3.16xlarge',
                        sagemaker_session=sagemaker_session,
                        # Training using SMDataParallel Distributed Training Framework
                        distribution={'smdistributed':{
                                            'dataparallel':{
                                                    'enabled': True
                                             }
                                      }}
                        )

estimator.fit()

# ## 推論
#
# 学習されたモデルができたので、モデルをホストするエンドポイントをデプロイします。エンドポイントをデプロイした後、推論リクエストでテストすることができます。

model_data = estimator.model_data
print("Using this model: {}".format(model_data))

# ### モデルオブジェクトの作成
# SageMaker SDK の TensorFlowModel を使ってモデルオブジェクトを定義し、 estimator と entry_point からモデルを渡します。関数はモデルをロードし、GPU が利用可能であれば GPU を使用するように設定します。

# +
import sagemaker
role = sagemaker.get_execution_role()

from sagemaker.tensorflow import TensorFlowModel
model = TensorFlowModel(model_data=model_data, role=role, framework_version='2.3')
# -

# ### モデルをエンドポイントにデプロイ
# model.deploy 関数を使用して predictor を作成します。オプションでインスタンス数とインスタンスタイプの両方を変更することができます。

predictor = model.deploy(initial_instance_count=1, instance_type='ml.m4.xlarge')

# +
# Download the test set
import tensorflow as tf

dataset = tf.keras.datasets.mnist.load_data(
    path='mnist.npz'
)

_, (test_imgs, test_labels) = dataset

# Randomly select 16 images from the test images
import numpy as np
import random

mask = random.sample(range(0, len(test_imgs)), 16)
mask = np.array(mask, dtype=np.int8)
samples = test_imgs[mask]

# Inspect sample images
import matplotlib.pyplot as plt
# %matplotlib inline

fig, ax = plt.subplots(nrows=2, ncols=8, sharex=True, sharey=True)
for i, row in enumerate(ax):
    for j, col in enumerate(row):
        col.imshow(samples[8*i+j].reshape(28, 28))

# +
# Send the samples to the endpoint for inference
samples = np.expand_dims(samples, axis=3)
outputs = predictor.predict(samples)['predictions']
outputs = np.array(outputs, dtype=np.float32)

print("Predictions: ")
print(np.argmax(outputs, axis=1))

print("Ground Truth: ")
print(test_labels[mask])
# -

# ## エンドポイントの削除
# エンドポイント への課金を止めるために、これ以上推論を試したり、エンドポイントを使う予定がない場合は、エンドポイントを削除してください。

predictor.delete_endpoint()



StarcoderdataJupyterScriptsDedupFiltered