Amazon SageMaker スクリプトモードでの TensorFlow Eager Execution の使用

このブログ記事では、TensorFlow の Eager Execution モードでモデルを訓練するために Amazon SageMaker のスクリプトモードを使用する方法について説明します。Eager Execution は TensorFlow の未来であり、現在最新の TensorFlow 1.x バージョンでオプションとして利用可能ではあるものの、TensorFlow 2 ではデフォルトモードになる予定です。今回は、スクリプトモードと Eager Execution について簡単に説明してから、代表的な回帰タスクシナリオをご紹介します。次に、スクリプトモードと Eager Execution を併用して、このタスクを解決するワークフローについて説明します。このブログ記事のためのノートブックと関連コードは GitHub から入手することができます。それでは、スクリプトモードから見ていきましょう。

Amazon SageMaker のスクリプトモード

Amazon SageMaker は、TensorFlow などの人気の高い機械学習 (ML) および深層学習フレームワークの一部を使ったモデルの訓練とデプロイメントを容易にする API と事前構築されたコンテナを提供します。Amazon SageMaker は、コンテナの構築や基盤となるインフラストラクチャの管理について心配することなく、カスタム TensorFlow コードを使用したモデルを訓練してデプロイするために使用できます。Amazon SageMaker の Python SDK TensorFlow Estimator、および Amazon SageMaker のオープンソース TensorFlow コンテナは、TensorFlow スクリプトを記述して、それを Amazon SageMaker で実行することを容易にします。これらの機能を活用するための望ましい方法が、スクリプトモードの使用です。

Amazon SageMaker のスクリプトモードは、AWS re:Invent 2018 でローンチされました。これは、特定の関数の定義済みインターフェイスと TensorFlow Estimator API を中心としたトレーニングコードの構造化を必要とする以前のレガシーモードに代わるものです。TensorFlow バージョン 1.11 からは、SageMaker 外で使用するものと同じ種類のトレーニングスクリプトで TensorFlow モデルを訓練するために、Amazon SageMaker の事前構築された TensorFlow コンテナでスクリプトモードを使用できます。スクリプトモードのコードには、特定の Amazon SageMaker 定義のインターフェイスへの適合も、特定の TensorFlow API の使用も必要ありません。

スクリプトモードのトレーニングスクリプトは Amazon SageMaker 外で使用するトレーニングスクリプトとよく似ていますが、独自に設定する様々な環境変数を介して Amazon SageMaker トレーニング環境の有用なプロパティにもアクセスできます。例えば、データセットの場所 (ローカルまたは Amazon S3 内)、およびアルゴリズムのハイパーパラメーターを指定するためにこれらの環境変数を使用します。以下のコードスニペットにあるように、コードが Python で記述されている場合は、通常実際のトレーニングを行うコードがメインガード (if __name__ == “__main__”) に置かれます。これは Amazon SageMaker がこのスクリプトをインポートするからです。このメインガードは、Amazon SageMaker でコードを実行する準備が整うまで、コードが実行されるのを阻止します。


if __name__ == "__main__":
        
    args, _ = parse_args()
    
    x_train, y_train = get_train_data(args.train)
    x_test, y_test = get_test_data(args.test)
    
    device = '/cpu:0' 
    print(device)
    batch_size = args.batch_size
    epochs = args.epochs
    print('batch_size = {}, epochs = {}'.format(batch_size, epochs))

    with tf.device(device):
        
        model = get_model()
        optimizer = tf.train.GradientDescentOptimizer(0.1)
        model.compile(optimizer=optimizer, loss='mse')    
        model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
                  validation_data=(x_test, y_test))

        # evaluate on test set
        scores = model.evaluate(x_test, y_test, batch_size, verbose=2)
        print("Test MSE :", scores)

        # save checkpoint for locally loading in notebook
        saver = tfe.Saver(model.variables)
        saver.save(args.model_dir + '/weights.ckpt')
        # create a separate SavedModel for deployment to a SageMaker endpoint with TensorFlow Serving
        tf.contrib.saved_model.save_keras_model(model, args.model_dir)

典型的なスクリプトモードスクリプトのフローでは、最初にコマンドライン引数が取得され、その後データがロードされて、モデルがセットアップされます。その次は実際のトレーニングです。これは、tf.keras によって提供される便利なメソッド、またはユーザー定義のトレーニングループを使って行われます。保存されたモデルはコンテナの /opt/ml/model ディレクトリに入れられ、トレーニング完了時におけるコンテナのティアダウン前に、ここから自動で Amazon S3 にアップロードされます。

TensorFlow の Eager Execution

Eager Execution は TensorFlow の未来であり、大きなパラダイムシフトです。TensorFlow の元々のグラフモードの、より直観的で動的な代替機能として最近導入された Eager Execution は、TensorFlow 2 のデフォルトモードになります。

Eager Execution のインターフェイスは無条件なもので、演算は、静的な計算グラフの構築に使用されるのではなく、直ちに実行されます。Eager Execution の利点には、自然な制御フローと少ないボイラープレートによるより直観的なインターフェイス、シンプル化されたデバッグ、および動的モデルとほぼすべての利用可能な TensorFlow 演算の対するサポートが含まれます。Eager Execution とグラフモードのもうひとつの重要な違いは、グラフモードでは変数などのプログラム状態,がグローバルに保存され、状態オブジェクトの寿命が tf.Session オブジェクトによって管理される一方、Eager Execution では、状態オブジェクトの寿命が、それらに対応する Python オブジェクトの寿命によって決定されます。これにより、コードが機能する方法の推測と、そのデバッグが簡単になります。

これらの利点に加えて、Eager Execution は tf.keras API とも連動し、迅速なプロトタイピングをさらに容易にします。tf.keras が使用される場合、モデルは tf.keras functional API、または tf.Keras.Model からのサブクラスを使って構築できます。tf.keras でのモデルのセットアップ後は、モデルをコンパイルし、fit メソッドを呼び出して訓練を行い、テストセットで評価して、保存することができます。tf.keras を使用しない場合は、独自のトレーニングループを定義してから、tf.GradientTape を使って後の自動微分のために演算を記録します。tf.keras を使用するかしないかにかかわらず、Amazon SageMaker の事前構築された TensorFlow コンテナで Eager Execution を使用できるようになりました。これはレガシーモードでは不可能でしたが、今ではスクリプトモードで有効化できます。

ワークフローの初期ステップ: データ処理とローカルモード

スクリプトモードで Eager Execution がどのように機能するかを説明するため、Amazon SageMaker 内の比較的完全なワークフローをお見せすることに焦点を当てていきます。このワークフローには、単純な回帰タスクという状況下における、ローカルおよびホスト型のトレーニング、ならびに推論が含まれます。このタスクには、良く知られている公の Boston Housing Dataset に基づいた住宅価格の予測が関与します。このデータセットには、平均部屋数、放射高速道路へのアクセス性、チャールズ川への近さなどを含む、ボストン近郊の町における住宅ストックに適用される 13 の特性があります。このブログ記事のステップを一緒に進めていくためにも、Amazon SageMaker ノートブックインスタンスをセットアップすることをお勧めします。すでにお持ちでない場合は、Amazon SageMaker 開発者ガイドにある手順を参照してください。ノートブックと関連コードは、このブログ記事のための GitHub リポジトリからアップロードできます。

データの前処理を行ってトレーニングスクリプトを作成した後の次のステップは、コードが期待通りに機能することの確認です。例えば、ごく少数のエポックのためにモデルを訓練する、または完全なデータセットではなくデータセットの小規模サンプルだけでモデルを訓練するのかもしれません。これを実行するための便利な方法は、Amazon SageMaker のローカルモードトレーニングを使用することです。ローカルモードで訓練するには、ノートブックインスタンスに Docker Compose または NVIDIA-Docker-Compose (GPU 向け) がインストールされている必要があります。サンプルコードには、これをチェックして、インストールされていないソフトウェアがあればそれらをインストールするセットアップシェルスクリプトがあります。

以下のコードスニペットは、TensorFlow Estimator をセットアップしてから少数のエポックでのトレーニングジョブを開始して、コードが機能していることを確認する方法を示しています。Estimator の主要パラメーターのひとつは train_instance_type で、これはトレーニングが実行されるハードウェアの種類です。ローカルモードについては、このパラメーターを単に「local」に設定して CPU でローカルモードトレーニングを呼び出す、またはインスタンスに GPU がある場合は「local_gpu」に設定します。他の主要パラメーターには、ディクショナリーとして渡されるアルゴリズムのハイパーパラメーター、およびスクリプトモードを使用していることを示すブーリアン型パラメーターがあります。

import sagemaker
from sagemaker.tensorflow import TensorFlow

model_dir = '/opt/ml/model'
train_instance_type = 'local'
hyperparameters = {'epochs': 10, 'batch_size': 128}
local_estimator = TensorFlow(entry_point='train.py',
                       model_dir=model_dir,
                       train_instance_type=train_instance_type,
                       train_instance_count=1,
                       hyperparameters=hyperparameters,
                       role=sagemaker.get_execution_role(),
                       base_job_name='tf-eager-scriptmode-bostonhousing',
                       framework_version='1.12.0',
                       py_version='py3',
                       script_mode=True)

inputs = {'train': f'file://{train_dir}',
          'test': f'file://{test_dir}'}

local_estimator.fit(inputs)

トレーニングジョブを開始するには、local_estimator.fit(inputs) を呼び出します。ここでの inputs は、キーと名前付きチャネルにデータセットの場所をポイントする値があるディクショナリーです。 local_estimator.fit(inputs) 呼び出しは、TensorFlow for Python 3、CPU バージョンの事前構築された TensorFlow コンテナをノートブックインスタンスにローカルにダウンロードします。これはその後、Amazon SageMaker のトレーニングジョブをシミュレートします。トレーニングが開始されると、TensorFlow コンテナが train.py スクリプトを実行し、ハイパーパラメーターをコマンドラインスクリプト引数として渡します。スクリプトが機能しているかどうかは、各トレーニングエポックのメトリクスを含む、ノートブックセル内の出力であるログを表示することで確認できます。

ローカルモードでコードが機能していることを確認した後は、取得してどこにでも (ノートブックインスタンスも含まれます) ロードできる、Amazon S3 に保存されたモデルチェックポイントもできています。さらなるサニティーチェックとして、今度はモデルを使って予測を行い、それらをテストセットと比較します。これを Boston Housing Dataset で行う場合は、住宅価格の値が 1,000 USD 単位であることに留意してください。(実際の時価が今日の大都市における住宅価格と比べて比較的低いことを不思議に思われている場合のために言うと、このデータセットを引用する文書は当初 1978 年に発行されたものです。) コードが機能していることを確認したら、ホスト型トレーニングに進みましょう。

Amazon SageMaker におけるホスト型トレーニング

ホスト型トレーニングは、完全なデータセットにおける完全なトレーニング (特に大規模な分散型トレーニング) の実行に望ましい方法です。ローカルモードトレーニングを実行したとき、データへのアクセスはローカルディレクトリから行われました。Amazon S3 は、すべてのデータを一か所にまとめておきたい場合に、ローカルモード用のトレーニングデータの保持にも使用できることを覚えておいてください。しかし、ホスト型トレーニングを開始する前に、ノートブックに示されている通り、データを Amazon S3 バケットにアップロードする必要があります。

データをアップロードしたら、Amazon SageMaker Estimator オブジェクトをセットアップする準備が整います。これは、(1) train_instance_type がローカルモードの「local」ではなく、特定のインスタンスタイプに設定されていること、および (2) fit 呼び出しに対する inputs 引数が Amazon S3 の場所に設定されていることを除き、ローカルモードの Estimator に似ています。また、本格的なトレーニングを実行する準備が整っていることから、エポック数も増やされています。これらの変更を行ったら、fit メソッドを再度呼び出すだけで実際のホスト型トレーニングを開始できます。

train_instance_type = 'ml.c4.xlarge'
hyperparameters = {'epochs': 30, 'batch_size': 128}

estimator = TensorFlow(entry_point='train.py',
                       model_dir=model_dir,
                       train_instance_type=train_instance_type,
                       train_instance_count=1,
                       hyperparameters=hyperparameters,
                       role=sagemaker.get_execution_role(),
                       base_job_name='tf-eager-scriptmode-bostonhousing',
                       framework_version='1.12.0',
                       py_version='py3',
                       script_mode=True)

estimator.fit(inputs)

predictor = estimator.deploy(initial_instance_count=1,instance_type='ml.m4.xlarge')
results = predictor.predict(x_test[:10])['predictions']

ローカルモードトレーニングと同様に、ホスト型トレーニングも、取得してロードできる Amazon S3 に保存されたモデルを生成します。その後、予測を行って、それらをテストセットと比較できます。これは、Amazon SageMaker のモジュール性も実証しています。モデルを Amazon SageMaker で訓練したため、このモデルは Amazon SageMaker から持ち出して、どこででも実行することができるようになりました。

その代わりに、Amazon SageMaker によってホストされるエンドポイントの機能を使ってモデルをデプロイすることもできます。これを TensorFlow で行うには、TensorFlow Serving で必須とされている通り、モデルをモデルチェックポイント形式ではなく、TensorFlow SavedModel 形式で保存する必要があります。上記にある最後のコードスニペットにあるように、ホストされたエンドポイントへのデプロイメントは、Estimator の deploy メソッドを呼び出すことによって、たった 1 行のコードで実現できます。

結論

Amazon SageMaker の目的のひとつは、データサイエンティストと開発者が ML モデルを迅速かつ簡単に構築、訓練、デプロイできるようにすることです。スクリプトモードと TensorFlow Eager Execution モードを組み合わせると、多種多様なデータサイエンスプロジェクトのために使用できる、大規模なトレーニングとデプロイメントに対する迅速なプロトタイピングのためのワークフローを簡単にセットアップできます。スクリプトモードは、TensorFlow の元の静的計算グラフモードを使いたい場合にも使用できます。ご自由に選択してください。これは、Amazon SageMaker が提供する数多くの柔軟なオプションのひとつに過ぎません。

著者について

Brent Rabowsky は AWS のデータサイエンスに焦点を当てており、彼の専門知識をAWS のお客様が独自のデータサイエンスプロジェクトを実行できるように活用しています。

Amazon Web Services ブログ