エンジニア目線で始める Amazon SageMaker Training ④-1 TensorFlow のコードを SageMaker Training 用へ書き換える

## 記事一覧

https://qiita.com/kazuneet/items/795e561efce8c874d115

https://qiita.com/kazuneet/items/294357a2214ac733f6f1

https://qiita.com/kazuneet/items/caa887826e2c33309e0b

https://qiita.com/kazuneet/items/bba523109f5fac0f410b

## 既存 TensorFlow のトレーニングコードを SageMaker Training 用に書き換える

[①の記事](https://qiita.com/kazuneet/items/795e561efce8c874d115)で SageMaker Training はどんなものでどうやって動かすのか、機械学習のコードを一切書かずに説明しました。

一方、実際どうやって機械学習を動かせばいいのかいまいちピンと来ない方もいらっしゃると思います。

この記事ではどうやれば既存の機械学習のトレーニングコードを SageMaker Training で動かせるのか、にフォーカスして説明していきます。

※不本意ながらついに機械学習のコードを使います
※例によってコードはすべて GitHub にあるので、Cloneしてお使いください。

https://github.com/kazuhitogo/sagemaker-training-tutorial

※この記事は TensorFlow(Keras) を使っていますので、PyTorch をお使いの方は以下リンクをご参照ください。

https://qiita.com/kazuneet/items/bba523109f5fac0f410b

## トレーニングコード

TensorFlow(Keras) でこんなコードを用意しました。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow
import tensorflow as tf

# データロード
(x_train,y_train),(x_valid,y_valid) = tf.keras.datasets.cifar10.load_data()
x_train = x_train/255
x_valid = x_valid/255

# モデル定義
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Conv2D(filters=16,kernel_size=(3,3),padding="same",activation="relu",input_shape=(x_train.shape[1],x_train.shape[2],x_train.shape[3])))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(10,activation="softmax"))

# トレーニング
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.00001),metrics=['accuracy'],loss="sparse_categorical_crossentropy")
model.fit(x_train,y_train,batch_size=4,epochs=2,validation_data=(x_valid,y_valid))

# モデルの保存
model.save('./1')
```

データはみんな大好き MNI… ではなく意味もなく CIFER-10 を使ってみました。
CIFER-10 のデータセットを TensorFlow から呼び出しています。

モデルはフィルター数 16 の畳み込みを一回だけ入れ、活性化関数に ReLU を指定して、最後に全結合層を入れて活性化関数に Softmax を入れて分類問題の学習をしているだけの、ものすごく簡単な構造です。

最後にモデルを SavedModel フォーマットで保存しています。

このトレーニングコードは Jupyter でそのまま動きます。

```text:jupyterの出力
(略)
Epoch 1/2
12500/12500 [==============================] - 26s 2ms/step - loss: 1.9984 - accuracy: 0.3121 - val_loss: 1.8459 - val_accuracy: 0.3635
Epoch 2/2
12500/12500 [==============================] - 25s 2ms/step - loss: 1.8035 - accuracy: 0.3804 - val_loss: 1.7620 - val_accuracy: 0.3933
(略)
INFO:tensorflow:Assets written to: ./1/assets
(略)
```


## 1. SageMaker Training で動かす

実は特にコードを変更しなくても動きます。
全く同じコードを、`src/4-1-1/train.py` に用意して以下のコードを実行します。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
import sagemaker
from sagemaker.tensorflow import TensorFlow

estimator = TensorFlow(
    entry_point='./src/4-1-1/train.py',
    py_version='py38', 
    framework_version='2.6.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit()
```

実行すると、処理は正常終了します。

```text:実行結果
(略)
INFO:tensorflow:Assets written to: ./1/assets
(略)
```

というわけで、SageMaker Training 用に特に書き換える必要もなく処理は動くのですが、SageMaker Training は揮発する(処理が終わるとインスタンスが消失する）ので、せっかくトレーニングしたモデルも揮発してしまいます。しかし、SageMaker Training では、①の記事の通り、トレーニングの成果物を指定ディレクトリに保存すれば Amazon S3 に転送してくれます。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E5%87%A6%E7%90%86%E7%B5%90%E6%9E%9C%E3%82%92%E4%BF%9D%E5%AD%98%E3%81%99%E3%82%8B%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E3%81%A7%E8%A8%80%E3%81%86%E3%81%A8%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E4%BF%9D%E5%AD%98

そこを実装してみましょう。

## 2.トレーニングが完了したらモデルを S3 に保存する

早速トレーニングコードを修正してみましょう。

```diff:./src/4-1-2/train.pyと./src/4-1-1/train.pyのdiff
(略)

+ import os

(略)


+ model_dir = os.environ.get('SM_MODEL_DIR')
- model.save('./1')
+ model.save(os.path.join(model_dir,'./1'))
```

①の記事の通り、`SM_MODEL_DIR` という環境変数で指定されているディレクトリに保存すると、処理完了時に Amazon S3 にデータが転送されるため、環境変数の取得と保存先の変更を行いました。

このトレーニングコードを SageMaker で動かしてみましょう。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
import sagemaker
from sagemaker.tensorflow import TensorFlow

estimator = TensorFlow(
    entry_point='./src/4-1-2/train.py',
    py_version='py38', 
    framework_version='2.6.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit()
```

無事トレーニングが終了しますが、Amazon S3 にモデルが転送されたか、モデルを読み込んで確認してみましょう。

```python./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import tensorflow as tf
model = tf.keras.models.load_model('./1')
model.summary()
```

```text:実行結果
(略)
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 32, 32, 16)        448       
_________________________________________________________________
flatten (Flatten)            (None, 16384)             0         
_________________________________________________________________
dense (Dense)                (None, 10)                163850    
=================================================================
Total params: 164,298
Trainable params: 164,298
Non-trainable params: 0
_________________________________________________________________
```

無事トレーニング完了後にモデルが S3 に転送され、使えることが確認できました。

## 3. トレーニングデータをコードの外に出す
機械学習のトレーニングではデータが変わることが多いですが、固定データの呼び出しがコードの中に埋め込まれていると毎回コードを変更する必要が出てきてしまうため、使い勝手が良くないです。
以下記事で紹介した通り、学習データを外から与えて(S3に配置したデータからトレーニングインスタンスに転送）やることで、トレーニングジョブ呼び出し時に使用するデータを決定し、トレーニングコードは与えられたデータを読み込むだけの実装に変更してみましょう。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E7%94%A8%E6%84%8F%E3%81%97%E3%81%9F%E3%83%87%E3%83%BC%E3%82%BF%E3%82%92-sagemaker-training-%E5%86%85%E3%81%A7%E4%BD%BF%E3%81%86%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E3%81%A7%E8%A8%80%E3%81%86%E3%81%A8%E5%AD%A6%E7%BF%92%E3%83%87%E3%83%BC%E3%82%BF%E3%81%AE%E5%88%A9%E7%94%A8

まずはデータを S3 に保存します。今回は numpy array をバイナリファイル(npy形式)に保存します。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
!mkdir -p ./train
!mkdir -p ./valid

import tensorflow as tf
import numpy as np
import sagemaker

(x_train,y_train),(x_valid,y_valid) = tf.keras.datasets.cifar10.load_data()
x_train = x_train/255
x_valid = x_valid/255

np.save('./train/x_train.npy',x_train)
np.save('./train/y_train.npy',y_train)
np.save('./valid/x_valid.npy',x_valid)
np.save('./valid/y_valid.npy',y_valid)
train_s3_uri = sagemaker.session.Session().upload_data(path='./train/', bucket=sagemaker.session.Session().default_bucket(), key_prefix='training/4-1/train')
valid_s3_uri = sagemaker.session.Session().upload_data(path='./valid/', bucket=sagemaker.session.Session().default_bucket(), key_prefix='training/4-1/valid')
```

これで `train_s3_uri` と  `valid_s3_uri` が指す S3 の URI に npy ファイル(x_train.npy, y_train.npy, x_valid.npy, y_valid.npy) をアップロードすることができました。

次にトレーニングコードを書き換えます。
後述しますが、`x_train.npy` と `y_train.npy` は 環境変数 `SM_CHANNEL_TRAIN` の値が示すディレクトリに、`x_valid.npy` と `y_valid.npy` は 環境変数 `SM_CHANNEL_VALID` の値が示すディレクトリに転送されるようにジョブを設定しますので、それぞれのディレクトリに各ファイルが存在する前提でコードを記載します。

```diff:./src/4-1-3/train.pyと./src/4-1-2/train.pyのdiff
(略)

+ import numpy as np

(略) 


- (x_train,y_train),(x_valid,y_valid) = tf.keras.datasets.cifar10.load_data()
- x_train = x_train/255
- x_valid = x_valid/255
+ train_dir = os.environ.get('SM_CHANNEL_TRAIN')
+ valid_dir = os.environ.get('SM_CHANNEL_VALID')
+ x_train = np.load(os.path.join(train_dir,'x_train.npy'))
+ y_train = np.load(os.path.join(train_dir,'y_train.npy'))
+ x_valid = np.load(os.path.join(valid_dir,'x_valid.npy'))
+ y_valid = np.load(os.path.join(valid_dir,'y_valid.npy'))
```

これで書き換え終了です。トレーニングジョブを起動しましょう。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
import sagemaker
from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
    entry_point='./src/4-1-3/train.py',
    py_version='py38', 
    framework_version='2.6.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit({
    'train':train_s3_uri,
    'valid':valid_s3_uri,
})
```

先程、`SM_CHANNEL_TRAIN` や `SM_CHANNEL_VALID` の値が示すディレクトリにファイルが転送されるように設定すると記載しました。
①の記事の通り `fit()` メソッドの引数に S3 の URI を指定するとトレーニングインスタンスにデータが転送されますが、転送先の指定を行うには引数を辞書形式で渡します。
辞書のキーが環境変数に設定され SM_CHANNEL_{キー} という形で引き渡されます。上のコードの例であれば、`train` というキーで指定された場合、`SM_CHANNEL_TRAIN` という環境変数の変数名に、`valid` というキーで指定された場合は`SM_CHANNEL_VALID`という環境変数の変数名で受け渡され、辞書の値で指定された S3 の URI で示されたデータが環境変数の値で示されるディレクトリに保存されます。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E8%A4%87%E6%95%B0%E3%81%AE-prefix-%E3%81%AB%E3%81%82%E3%82%8B%E3%83%87%E3%83%BC%E3%82%BF%E3%82%92%E4%BD%BF%E3%81%86-training-job

処理が無事完了したら、念の為モデルを確認しましょう。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!rm model.tar.gz
!rm -rf ./1
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import tensorflow as tf
model = tf.keras.models.load_model('./1')
model.summary()
```

```text:実行結果
(略)
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 32, 32, 16)        448       
_________________________________________________________________
flatten (Flatten)            (None, 16384)             0         
_________________________________________________________________
dense (Dense)                (None, 10)                163850    
=================================================================
Total params: 164,298
Trainable params: 164,298
Non-trainable params: 0
_________________________________________________________________
```

しっかりモデルが出来上がったことがわかります。


## 4. ハイパーパラメータを外から渡す
だいぶ SageMaker っぽい使い方になってきましたが、もう少し頑張ってみましょう。
モデルを試行錯誤する上でハイパーパラメータの値は変更しやすいと嬉しいです。そのハイパーパラメータ変更のたびに都度コードを変更するのは効率が良くないので、トレーニングジョブ起動時に設定するようにしましょう。また、設定しなくても動くようにデフォルト値も設定しておきましょう。
このケースでは、畳み込みのフィルター数と、バッチサイズ、エポック数と、学習率を外部から与えてみます。

トレーニングコードをこのように修正しました。
```diff:./src/4-1-4/train.pyと./src/4-1-3/train.pyのdiff
(略)

+ # ハイパーパラメータ取得
+ hps = json.loads(os.environ.get('SM_HPS'))
+ hps.setdefault('batch_size', 4)
+ hps.setdefault('epochs', 2)
+ hps.setdefault('filters', 16)
+ hps.setdefault('learning_rate', 0.00001)
+ print(hps)

(略)


- model.add(tf.keras.layers.Conv2D(filters=16,kernel_size=(3,3),padding="same",activation="relu",input_shape=(x_train.shape[1],x_train.shape[2],x_train.shape[3])))
+ model.add(tf.keras.layers.Conv2D(filters=hps['filters'],kernel_size=(3,3),padding="same",activation="relu",input_shape=(x_train.shape[1],x_train.shape[2],x_train.shape[3])))

(略)


- model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.00001),metrics=['accuracy'],loss="sparse_categorical_crossentropy")
+ model.fit(x_train,y_train,batch_size=hps['batch_size'],epochs=hps['epochs'],validation_data=(x_valid,y_valid))

```

ハイパーパラメータを取得するやり方はコマンドライン引数から取得する方法(`argparse` を使ってパースする)と環境変数から取得する方法がありますが、
個人の好みで環境変数から取得してみました。併せて`setdefault()`を使って定義されていない場合のデフォルト値をセットしています。

さてトレーニングコードを実行してみます。ハイパーパラメータが外から与えられているか確認するため、デフォルト値とは別のハイパーパラメータを与えます。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
import sagemaker
from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
    entry_point='./src/4-1-4/train.py',
    py_version='py38', 
    framework_version='2.6.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role(),
    hyperparameters={
        'filters':8,
        'epochs':1,
        'batch-size':'16',
        'learning_rate' : 0.001
    }
)
estimator.fit({
    'train':train_s3_uri,
    'valid':valid_s3_uri,
})
```

```text:実行結果
(略)
{'batch_size': 16, 'epochs': 1, 'filters': 8, 'learning_rate': 0.001, 'model_dir': 's3://sagemaker-{REGION}-{ACCOUNT_ID}/tensorflow-training-{YYYY-MM-DD-HH-MI-SS-mmm}/model'}
(略)
INFO:tensorflow:Assets written to: /opt/ml/model/./1/assets
2022-03-09 02:49:19,499 sagemaker-training-toolkit INFO     Reporting training SUCCESS

2022-03-09 02:49:36 Uploading - Uploading generated training model
2022-03-09 02:49:36 Completed - Training job completed
(略)
```

無事完了しました。ハイパーパラメータが効いているか、モデルをロードしてフィルター数だけ確認しておきましょう。

```python:./4_1_rewriting_traing_code_for_sagemaker_tensorflow.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!rm model.tar.gz
!rm -rf ./1
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import tensorflow as tf
model = tf.keras.models.load_model('./1')
model.summary()
```

```text:実行結果
(略)
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 32, 32, 8)         224       
_________________________________________________________________
flatten (Flatten)            (None, 8192)              0         
_________________________________________________________________
dense (Dense)                (None, 10)                81930     
=================================================================
Total params: 82,154
Trainable params: 82,154
Non-trainable params: 0
_________________________________________________________________

```

無事`(None, 32, 32, 8) `とフィルター数が8になっていることを確認できました。

## 終わりに
TensorFlow(Keras) の機械学習コードを SageMaker で動かすのに適したコードに書き換えました。
ここまでの機能を利用することで、SageMaker のハイパーパラメータチューニングの機能も簡単に利用できるようになるので、ぜひこの使い方を覚えましょう。

https://aws.amazon.com/jp/blogs/news/sagemaker-automatic-model-tuning/
https://sagemaker.readthedocs.io/en/stable/api/training/tuner.html