Amazon Web Services ブログ
Amazon SageMaker で簡単に Keras を使う方法
Amazon SageMaker は、すべての開発者とデータサイエンティストに機械学習モデルの構築、トレーニング、デプロイの手段を提供する AWS のマネージドサービスです。SageMaker は深層学習の主要なフレームワークをサポートしており、TensorFlow、Apache MXNet、PyTorch、Chainer などを用いてモデルを記述することができます。また、TensorFlow や MXNet をバックエンドとした Keras の高レベルな API を用いることで、モデルを簡単にすばやく作成することもできます。
これまで、Keras で書いたコードを SageMaker 上で動かす方法について、多くのお客様からご質問を頂いておりました。2018年12月に、SageMaker の TensorFlow ならびに MXNet のコンテナに、それぞれのフレームワークをバックエンドとする Keras がデフォルトでインストールされました。また両コンテナでは Script Mode をサポートしているため、SageMaker 外で開発した Keras のコードに、わずかな修正を加えるだけで動かすことができるようになりました。ここでは Keras 公式サンプルコードの mnist_cnn.py をなるべくそのまま利用して、SageMakerのTensorFlowとMXNetコンテナで実行する方法をご説明します。
TensorFlow Backend での Keras の使い方
トレーニングスクリプトの修正
AWS のマネージドコンソールから SageMaker ノートブックインスタンス (Jupyter/JupyterLab) を開き、Keras のリポジトリをクローンします (このブログのようにノートブックインスタンスの作成時に関連付けることも可能です)。keras/examples/mnist_cnn.py の中で以下の3点を修正します:
- 学習からモデルの保存までを
train(args)関数として定義します。ここでは次の手順で読み込むargsでパラメータを受け取ります。S3 から取得したデータを読み込みトレーニングを行います。最後に、後の手順で定義するsave(model, model_dir)関数でモデルを保存します。import numpy as np def train(args): batch_size = args.batch_size epochs = args.epochs num_classes = args.num_classes train_dir = args.train # load data x_train = np.load(os.path.join(train_dir, 'train.npz'))['image'] y_train = np.load(os.path.join(train_dir, 'train.npz'))['label'] x_test = np.load(os.path.join(train_dir, 'test.npz'))['image'] y_test = np.load(os.path.join(train_dir, 'test.npz'))['label'] # the original training script starts here # input image dimensions img_rows, img_cols = 28, 28 if K.image_data_format() == 'channels_first': x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols) x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols) input_shape = (1, img_rows, img_cols) else: x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1) x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1) input_shape = (img_rows, img_cols, 1) ... score = model.evaluate(x_test, y_test, verbose=0) print('Test loss:', score[0]) print('Test accuracy:', score[1]) save(model, args.model_dir) import argparse, osを読み込んだ上で、main guard (if __name__ == '__main__':) の中で CreateTrainingJob API から渡されたパラメータをパースします。 コマンドライン引数および環境変数として、ハイパーパラメータと Amazon S3 上のトレーニングデータ・モデルの保存場所が取得できます。このパラメータを渡して、上で定義したtrain()関数を呼びます。if __name__ == '__main__': parser = argparse.ArgumentParser() # hyperparameters sent by the client are passed as command-line arguments to the script parser.add_argument('--batch-size', type=int, default=128) parser.add_argument('--num-classes', type=int, default=10) parser.add_argument('--epochs', type=int, default=12) # input data and model directories parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR']) parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING']) args, _ = parser.parse_known_args() train(args)train()の末尾で呼び出したsave(model, model_dir)関数では、トレーニング後 TensorFlow Serving でモデルをデプロイできるように TensorFlow のモデルを保存します。import tensorflow as tf from keras import backend as K def save(model, model_dir): sess = K.get_session() tf.saved_model.simple_save( sess, os.path.join(model_dir, 'model/1'), inputs={'inputs': model.input}, outputs={t.name: t for t in model.outputs})
TensorFlow Serving を使うと推論用の関数は実装する必要がないので、 Python スクリプトの中で必要な修正はこれだけです。修正済みのコードを今回は mnist_cnn_sagemaker_tensorflow.py という名前で保存することにします。
トレーニングジョブの作成
では次に、トレーニングジョブを呼び出すため Jupyter Notebook を作成します。カーネルは conda_tensorflow_p36 を選択して下さい。まずはデータセットを用意します。
import os
import keras
import numpy as np
from keras.datasets import mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
os.makedirs("./data", exist_ok = True)
np.savez('./data/train', image=x_train, label=y_train)
np.savez('./data/test', image=x_test, label=y_test)
import sagemaker
sagemaker_session = sagemaker.Session()
bucket_name = sagemaker_session.default_bucket()
input_data = sagemaker_session.upload_data(path='./data', bucket=bucket_name, key_prefix='dataset/mnist')
print('Training data is uploaded to: {}'.format(input_data))次に、先ほど書き換えたトレーニング用のスクリプトを entry_point に指定し、各種パラメータを指定して TensorFlow Estimator を初期化します。ここで py_version='py3', script_mode=True を指定することを忘れないで下さい。そして、上で S3 にアップロードした input_data を引数にして、fit() メソッドでトレーニングジョブを呼び出します。
from sagemaker.tensorflow import TensorFlow
from sagemaker import get_execution_role
role = get_execution_role()
estimator = TensorFlow(
entry_point = "./mnist_cnn_sagemaker_tensorflow.py",
role=role,
train_instance_count=1,
train_instance_type="ml.m5.xlarge",
framework_version="1.12.0",
py_version='py3',
script_mode=True,
hyperparameters={'batch-size': 64,
'num-classes': 10,
'epochs': 1})
estimator.fit(input_data)これだけです！あとはトレーニングジョブが作成されて最終的に S3 にモデルが出力されるのを待ちます。なお、Python スクリプトのコンソール出力はそのまま Notebook 上に表示されますし、後から Amazon CloudWatch Logs で確認することもできます。
推論エンドポイントへモデルのデプロイ
トレーニングに用いた Python スクリプトでは S3 にモデルを保存しました。これを TensorFlow Serving を使って推論エンドポイントにデプロイします。デプロイは以下のメソッドで簡単に実行できます。
predictor = estimator.deploy(instance_type='ml.m5.xlarge', initial_instance_count=1)しばらく待ってエンドポイントが立ち上がったら、実際に推論を行ってみましょう。
%matplotlib inline
import random
import matplotlib.pyplot as plt
num_samples = 5
indices = random.sample(range(x_test.shape[0] - 1), num_samples)
images, labels = x_test[indices]/255, y_test[indices]
for i in range(num_samples):
plt.subplot(1,num_samples,i+1)
plt.imshow(images[i].reshape(28, 28), cmap='gray')
plt.title(labels[i])
plt.axis('off')
prediction = predictor.predict(images.reshape(num_samples, 28, 28, 1))['predictions']
prediction = np.array(prediction)
predicted_label = prediction.argmax(axis=1)
print('The predicted labels are: {}'.format(predicted_label))このコードを実行し、以下のような推論結果が出力されれば成功です。[0 2 7 5 0] は、ランダムに選択された5枚の数字の画像に対する認識結果で、実行するたびに変化します。
なお、立ち上げたエンドポイントの削除は sagemaker.Session().delete_endpoint(predictor.endpoint) で行うことができます。
MXNet バックエンドで Keras を動かす場合
SageMaker の MXNet 環境には Fork された Keras が含まれているため、これを用いて MXNet バックエンドで Keras を動かすことができます。コマンドライン引数をパースして train(args) 関数として定義したコードを実行するという点は同じですが、モデルの保存・デプロイ方法が上記と少し異なります:
- 学習後は MXNet のモデルとして保存します。なお、MXNet のモデルを gluon.SymbolBlock.imports() で読み込む際は入力変数名を指定する必要があります。そのため、予め1層目のレイヤー (Conv2D) に
name='data'という名前を付けておきます。model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape, name='data'))その上で次のようにモデルを保存します。
def save(model, model_dir): keras.models.save_mxnet_model(model=model, prefix=os.path.join(model_dir, 'model')) - ホスティングのためモデルを読み込む関数
model_fn()を定義します。上で'data'という名前を付けた入力は'data_input1'に対応しています。def model_fn(model_dir): import mxnet as mx from mxnet import gluon net = gluon.nn.SymbolBlock.imports(os.path.join(model_dir, 'model-symbol.json'), ['/data_input1'], param_file=os.path.join(model_dir, 'model-0000.params'), ctx=mx.cpu()) return net
これを mnist_cnn_sagemaker_mxnet.py という名前で保存し、conda_mxnet_p36 カーネルのノートブックインスタンスから以下のコードでトレーニングジョブを呼び出します。
from sagemaker.mxnet import MXNet
from sagemaker import get_execution_role
role = get_execution_role()
estimator = MXNet("./mnist_cnn_sagemaker_mxnet.py",
role=role,
train_instance_count=1,
train_instance_type="ml.m5.xlarge",
framework_version="1.3.0",
py_version='py3',
hyperparameters={'batch-size': 64,
'num-classes': 10,
'epochs': 1})
estimator.fit(input_data)エンドポイントの作成が完了したら、推論を試すことができます。TensorFlow バックエンドの場合と同様ですが、SageMaker の MXNet コンテナはデフォルトで list を返すため以下のように結果を確認します。
prediction = predictor.predict(images.reshape(num_samples, 28, 28, 1))
prediction = np.array(prediction)
predicted_label = prediction.argmax(axis=1)
print('The predicted labels are: {}'.format(predicted_label))動作確認が完了し、必要のなくなったエンドポイントは sagemaker.Session().delete_endpoint(predictor.endpoint) で削除しましょう。
まとめ
本稿では Keras を利用してSageMaker上で機械学習のモデルを学習・デプロイする方法を、TensorFlow と MXNet バックエンド両方についてご紹介しました。Script Mode に対応した SageMaker Estimator を利用することで、Keras のコードを非常に簡単に SageMaker 上で実行することができます。上記以外の Keras の利用方法として、TensorFlow による Keras の実装 tf.keras を利用する方法もあります。SageMaker では今回ご紹介した構築・トレーニング・デプロイ以外にも、データのラベル付け、ハイパーパラメータ最適化、トレーニング済モデルのコンパイル、などの機能が簡単にお使い頂けます。
著者について
針原 佳貴 (Yoshitaka Haribara, Ph.D.) は AWS のソリューション アーキテクトです。趣味はドラムで休日は友達とバンドをして過ごしています。好きなバンドは Red Hot Chili Peppers、好きな AWS のサービスは Amazon SageMaker です。
鮫島 正樹 (Masaki Samejima, Ph.D.) は AWS の機械学習スペシャリスト ソリューション アーキテクトで、コンピュータビジョンや時系列解析などを得意にしています。