記事一覧
既存 TensorFlow のトレーニングコードを SageMaker Training 用に書き換える
①の記事で SageMaker Training はどんなものでどうやって動かすのか、機械学習のコードを一切書かずに説明しました。
一方、実際どうやって機械学習を動かせばいいのかいまいちピンと来ない方もいらっしゃると思います。
この記事ではどうやれば既存の機械学習のトレーニングコードを SageMaker Training で動かせるのか、にフォーカスして説明していきます。
※不本意ながらついに機械学習のコードを使います
※例によってコードはすべて GitHub にあるので、Cloneしてお使いください。
※この記事は TensorFlow(Keras) を使っていますので、PyTorch をお使いの方は以下リンクをご参照ください。
トレーニングコード
TensorFlow(Keras) でこんなコードを用意しました。
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 でそのまま動きます。
(略)
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 に用意して以下のコードを実行します。
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()
実行すると、処理は正常終了します。
(略)
INFO:tensorflow:Assets written to: ./1/assets
(略)
というわけで、SageMaker Training 用に特に書き換える必要もなく処理は動くのですが、SageMaker Training は揮発する(処理が終わるとインスタンスが消失する）ので、せっかくトレーニングしたモデルも揮発してしまいます。しかし、SageMaker Training では、①の記事の通り、トレーニングの成果物を指定ディレクトリに保存すれば Amazon S3 に転送してくれます。
そこを実装してみましょう。
2.トレーニングが完了したらモデルを S3 に保存する
早速トレーニングコードを修正してみましょう。
(略)
+ 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 で動かしてみましょう。
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 にモデルが転送されたか、モデルを読み込んで確認してみましょう。
# モデルの 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()
(略)
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に配置したデータからトレーニングインスタンスに転送）やることで、トレーニングジョブ呼び出し時に使用するデータを決定し、トレーニングコードは与えられたデータを読み込むだけの実装に変更してみましょう。
まずはデータを S3 に保存します。今回は numpy array をバイナリファイル(npy形式)に保存します。
!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 の値が示すディレクトリに転送されるようにジョブを設定しますので、それぞれのディレクトリに各ファイルが存在する前提でコードを記載します。
(略)
+ 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'))
これで書き換え終了です。トレーニングジョブを起動しましょう。
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 で示されたデータが環境変数の値で示されるディレクトリに保存されます。
処理が無事完了したら、念の為モデルを確認しましょう。
# モデルの 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()
(略)
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 っぽい使い方になってきましたが、もう少し頑張ってみましょう。
モデルを試行錯誤する上でハイパーパラメータの値は変更しやすいと嬉しいです。そのハイパーパラメータ変更のたびに都度コードを変更するのは効率が良くないので、トレーニングジョブ起動時に設定するようにしましょう。また、設定しなくても動くようにデフォルト値も設定しておきましょう。
このケースでは、畳み込みのフィルター数と、バッチサイズ、エポック数と、学習率を外部から与えてみます。
トレーニングコードをこのように修正しました。
(略)
+ # ハイパーパラメータ取得
+ 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()を使って定義されていない場合のデフォルト値をセットしています。
さてトレーニングコードを実行してみます。ハイパーパラメータが外から与えられているか確認するため、デフォルト値とは別のハイパーパラメータを与えます。
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,
})
(略)
{'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
(略)
無事完了しました。ハイパーパラメータが効いているか、モデルをロードしてフィルター数だけ確認しておきましょう。
# モデルの 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()
(略)
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