Amazon SageMaker での全体的な TensorFlow 2 ワークフローの作成

深層学習プロジェクトの全体的なライフサイクルを管理することには課題が伴います。特に、複数の別々なツールやサービスを使用する場合はそれが顕著です。たとえば、データの前処理、トレーニングと推論用コードのプロトタイピング、フルスケールでのモデルトレーニングとチューニング、モデルのデプロイ、そしてそれらすべてを稼働環境向けにオーケストレーションするためのワークフローの自動化などには、それぞれ異なるツールが使われるでしょう。ツールの入れ替えで手間がかかると、プロジェクトが遅延しコストは増大しがちです。今回のブログ記事では、Amazon SageMaker を使い、深層学習プロジェクトの全体的なライフサイクルを効率的に管理する方法をご紹介します。ここでは、サンプルコードのフレームワークとして TensorFlow 2 を使いますが、解説しているコンセプトは、他のフレームワークにも基本的に適用可能です。

また記事内では、関連したサンプルノートブックを取り上げていますが、これも、1 時間以内に実行が可能なものであり、解説しているすべての機能をお試しいただくことができます。詳細については、GitHub リポジトリをご参照ください。

Amazon SageMaker ワークフローの概要

TensorFlow 2 あるいは他のフレームワークを使用するデータサイエンスプロジェクトでは、データセットの取得、調査、前処理が、必ず作業の出発点となります。Amazon SageMaker ワークフローについて言えば、データの調査は一般的にノートブック内で行われます。こういったノートブックは、通常、就業時間帯のほとんどで実行されるため、比較的小型でさほどパワーがなく、料金も低いインスタンスタイプに適しています。

したがって、データセットが小さめな場合を除けば、ノートブックはフルスケールのデータ処理やモデルトレーニング、そして推論などに最適とは言えません。通常これらのタスクでは、相当量の並列コンピューティングリソースが必要となるため、ノートブックでの実行は現実的とは言えない訳です。一方で、サイズが最適でよりパワーがあり、これらのタスクが迅速に完了できる個別のインスタンスによるクラスターに、Amazon SageMaker のスピンアップ機能を使用すれば、かなり実践的でコスト効率も良くできます。この場合の料金は秒単位で発生し、各インスタンスはジョブの完了時点で Amazon SageMaker により自動的にシャットダウンされます。結果的に、Amazon SageMaker での通常のワークフローでの最も頻度の高い課金は、パワーがあり高価な GPU や高速化されたコンピューティングインスタンスではなく、データ調査やプロトタイピングを行う比較的低料金のノートブックで発生します。

プロトタイピングが完了すれば、自動化されたワークフローにより、作業はノートブックの上位へと移行されます。頑強かつ反復可能な手法でモデルのデプロイを行うワークフロー全体をオーケストレーションするためには、自動化されたパイプラインが不可欠です。そして、Amazon SageMaker では、このソリューションがネイティブに用意されています。この記事内の以降のセクションにおいて、プロジェクトライフサイクルでの各ステージで実装する、Amazon SageMaker のさまざまな機能をご紹介していきます。

Amazon SageMaker Processing によるデータ変換

Amazon SageMaker Processing は、ノートブックから分離されており適切なサイズで管理されたクラスターでの、大規模なデータセットの前処理に役立ちます。Amazon SageMaker Processing では、用意されたたサポートにより、Scikit-learn がすぐに使えます。また、他のコンテナ化されたテクノロジーもサポートしています。たとえば、Amazon SageMaker Processing の中で、機能変換のために一時的な Apache Spark クラスターを起動できます。

Amazon SageMaker Processing で Scikit-learn を使用するためには、標準的な Scikit-learn のコードを使った Python での、データ前処理スクリプトを記述する必要があります。このスクリプトにおける最小限の制約は、入力および出力のデータを、特定の場所に置く必要があるということのみです。Amazon SageMaker Processing は、入力データを Amazon Simple Storage Service (Amazon S3) から自動的に読み込み、ジョブが完了すると変換後のデータを Amazon S3 にアップロードします。

Amazon SageMaker Processing でのジョブを開始する前に、次に示すコードのように、SKLearnProcessor オブジェクトをインスタンス化します。このオブジェクト内で、ジョブが使用するインスタンスタイプと、インスタンスの数を指定します。

from sagemaker import get_execution_role
from sagemaker.sklearn.processing import SKLearnProcessor

sklearn_processor = SKLearnProcessor(framework_version='0.20.0',
                                     role=get_execution_role(),
                                     instance_type='ml.m5.xlarge',
                                     instance_count=2)

各クラスターインスタンスに対し、処理用に均等なデータファイルを振り分けるため、ディストリビューションタイプ ShardedByS3Key を ProcessingInput オブジェクトの中で指定します。これにより、インスタンスの数が n 個ある場合は、各インスタンスは指定した S3 バケットから 1/n 個のファイルを受け取るように定義されます。ステートレスなデータ変換用のインスタンスによる大規模なクラスターを簡単に作成できるという機能は、Amazon SageMaker Processing が提供する多くのメリットの中のほんの一例にすぎません。

from sagemaker.processing import ProcessingInput, ProcessingOutput
from time import gmtime, strftime 

processing_job_name = "tf-2-workflow-{}".format(strftime("%d-%H-%M-%S", gmtime()))
output_destination = 's3://{}/{}/data'.format(bucket, s3_prefix)

sklearn_processor.run(code='preprocessing.py',
                      job_name=processing_job_name,
                      inputs=[ProcessingInput(
                        source=raw_s3,
                        destination='/opt/ml/processing/input',
                        s3_data_distribution_type='ShardedByS3Key')],
                      outputs=[ProcessingOutput(output_name='train',
                                                destination='{}/train'.format(output_destination),
                                                source='/opt/ml/processing/train'),
                               ProcessingOutput(output_name='test',
                                                destination='{}/test'.format(output_destination),
                                                source='/opt/ml/processing/test')])

ローカルモードにおけるトレーニングと推論コードのプロトタイピング

データセットでトレーニングの準備ができたら、次に行うステップはトレーニング用コードのプロトタイピングです。TensorFlow 2 用として最も便利なワークフローは、トレーニング用の取り込みスクリプトを、Amazon SageMaker で構築済みの TensorFlow 2 コンテナにより作成することです。この機能は スクリプトモードと呼ばれており、Amazon SageMaker のローカルモードでのトレーニング機能と、シームレスに連携します。

ローカルモードは、ローカルのノートブック上でコードが機能することを確認するのに便利です。これは、Amazon SageMaker が管理する最適なサイズの個別クラスターで行う、ホスティングされたフルスケールのトレーニングに移行する前に使用します。ローカルモードにおいては、通常、少数のエポックに関し短時間だけ、可能であれば完全なデータセットのサンプルを使いトレーニングを行います。これによりコードが正確に機能することが確認でき、フルスケールでのトレーニング時間を無駄にすることも避けられます。インスタンスタイプについては、そのノートブックが GPU あるいは CPU のインスタンスに応じて、local_gpu か local のどちらかを指定します。

from sagemaker.tensorflow import TensorFlow

git_config = {'repo': 'https://github.com/aws-samples/amazon-sagemaker-script-mode',
              'branch': 'master'}

model_dir = '/opt/ml/model'
train_instance_type = 'local'
hyperparameters = {'epochs': 5, 'batch_size': 128, 'learning_rate': 0.01}
local_estimator = TensorFlow(git_config=git_config,
                             source_dir='tf-2-workflow/train_model',
                             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-2-workflow',
                             framework_version='2.1',
                             py_version='py3',
                             script_mode=True)

このローカルモードトレーニングは、フルスケールのトレーニングに移行する前のコードが機能しているのを確かめるのに便利ですが、同時に、推論コードのプロトタイピングを簡素化するための便利な手法でもあります。TensorFlow の SavedModel アーティファクト、もしくは Amazon S3 に保存してあるチェックポイントを取得し、テストのためにノートブックにロードするというのも、1 つの方法ではあります。しかし、ローカルモードのエンドポイントを使えば、これをより簡単に行えます。

ローカルモードのトレーニングジョブが提供するエスティメーターオブジェクトを使うことで、Amazon SageMaker の TensorFlow Serving コンテナを含むローカルモードエンドポイントに対し、モデルのデプロイが可能です。1 か所の例外を除けば、このコードは個別にホスティングされたエンドポイントでモデルをデプロイするコードと同じものです。こちらでは、ローカルのエスティメーターをデプロイするメソッドを呼び出してから、使用しているノートブックが GPU か CPU かにより、インスタンスタイプに local_gpu か local のどちらかを指定します。

local_predictor = local_estimator.deploy(initial_instance_count=1, instance_type='local')
local_results = local_predictor.predict(x_test[:10])['predictions']

ローカルモードを使用する前に、お使いのインスタンス上で、docker-compose もしくは nvidia-docker-compose (GPU インスタンスの場合) が実行可能であることを確認します。このブログ記事用の GitHub レポジトリには、このためのスクリプトをご用意しています。

自動モデルチューニング

プロトタイピングが完了した後は、Amazon SageMaker でホスティングするトレーニングと自動モデルチューニングを行っていきます。ホスト型トレーニングは、フルスケールトレーニング、特に大規模な分散型トレーニングの実行に適しています。ローカルモードと違いホスト型トレーニングでは、実際のトレーニングはノートブック上ではなく、Amazon SageMaker が管理する各マシンによる個別のクラスター上で実施されます。ホスト型トレーニング用のエスティメーターオブジェクトとローカルモードのエスティメーターの間の違いは、

データセットを Amazon S3、Amazon Elastic File System (Amazon EFS)、もしくは Amazon FSx for Lustre に置く必要があるということだけです。
トレーニング用のインスタンスタイプは、フルスケールのトレーニングに適したサイズの、Amazon SageMaker ML インスタンスタイプに設定されます。

ローカルモードのプロトタイピングにおいてトレーニングコードが機能することは確認済みなので、ホスト型トレーニングエスティメーターを修正して、多数のエポックと (仮にローカルモードでサンプルを使用した場合は) 完全なデータセットでのトレーニングすることができます。

とは言え、個々のホスト型トレーニングジョブを実行しながら、最適なモデルを見つけるためにハイパーパラメータを手動で調整する作業はかなり大変で、時間も浪費します。最適なハイパーパラメータの組合せの選択作業は、データセットとアルゴリズムに基づき変わってきます。アルゴリズムに調整すべき多くのハイパーパラメータがある場合とか、選択したハイパーパラメータに対し非常に繊細に反応する場合もあります。さらに、多くの場合は、フィットするモデルとハイパーパラメータの間の関係は直線的ではありません。自動モデルチューニングは、このチューニングプロセスをスピードアップします。これにより、異なるハイパーパラメータの組合せによる複数のトレーニングジョブを実行でき、パフォーマンスが最良になるモデルを特定できます。

次のサンプルコードに示すように、自動モデルチューニングを使うためには、まず、チューニングすべきハイパーパラメータを指定し、チューニングの幅と最適化の目標メトリクスを定義します。HyperparameterTuner オブジェクトが、これらをパラメータとして取得します。また、各チューニングジョブでは、そのジョブ内での最大トレーニング回数 (今回の場合は 15 回) の指定と、どの程度並列処理を行うか (今回は 1 回につき 5 つのジョブ) を指定する必要もあります。このパラメータ設定により、このチューニングジョブは、並列な 5 つのジョブを 3 回繰り返した後で終了します。デフォルトのベイズ最適化によるチューニング戦略では、今回のチューニング検索に対し以前のトレーニングジョブグループの結果が伝達されます。これにより、全体を並列処理するのではなく、いくつかのグループに並列ジョブを分割するのに役立ちます。ただ、より多くの並列ジョブを使用することでチューニングが速く終了するものの、チューニング検索の精度を犠牲にしやすいというトレードオフ関係も存在します。

from sagemaker.tuner import IntegerParameter, CategoricalParameter, ContinuousParameter, HyperparameterTuner

hyperparameter_ranges = {
  'learning_rate': ContinuousParameter(0.001, 0.2, scaling_type="Logarithmic"),
  'epochs': IntegerParameter(10, 50),
  'batch_size': IntegerParameter(64, 256),
}

metric_definitions = [{'Name': 'loss',
                       'Regex': ' loss: ([0-9\\.]+)'},
                     {'Name': 'val_loss',
                       'Regex': ' val_loss: ([0-9\\.]+)'}]

objective_metric_name = 'val_loss'
objective_type = 'Minimize'

tuner = HyperparameterTuner(estimator,
                            objective_metric_name,
                            hyperparameter_ranges,
                            metric_definitions,
                            max_jobs=15,
                            max_parallel_jobs=5,
                            objective_type=objective_type)

tuning_job_name = "tf-2-workflow-{}".format(strftime("%d-%H-%M-%S", gmtime()))
tuner.fit(inputs, job_name=tuning_job_name)

デプロイとワークフローの AWS Step Functions Data Science SDK による自動化

チューニングされた最良のモデルをデプロイするための便利なオプションは、Amazon SageMaker がホスティングするエンドポイントの利用することす。これによりリアルタイムの推論を実施できます (非同期でオフラインの推論用にバッチ変換ジョブも利用可能です。) 。このエンドポイントは、TensorFlow SavedModel を取得し、Amazon SageMaker TensorFlow Serving コンテナの中でそれをデプロイします。これは、HyperparameterTuner オブジェクトのデプロイメソッドを呼び出す、コード 1 行を記述するだけで実現できます。

tuning_predictor = tuner.deploy(initial_instance_count=1, instance_type='ml.m5.xlarge')

とは言え、プロトタイピングには便利なノートブックではありますが、これは通常、実稼働環境でのデプロイには使用されません。その代わりに、トレーニングとデプロイなど、複数のステップを含むパイプラインの実行にいは、ワークフローオーケストレーターが便利に使えます。たとえば、Amazon SageMaker における 1 つのパイプラインには、次のような 4 つのステップがあります。

モデルをトレーニングする。
使用するモデルアーティファクトを包含した Amazon SageMaker Model オブジェクトを作成する。
Amazon SageMaker エンドポイントの定義を作成し、そのモデルの保持方法 (インスタンスタイプやインスタンスの数など) を指定する。
トレーニング済みモデルを定義した Amazon SageMaker エンドポイントでデプロイする。

AWS Step Functions Data Science SDK は、このようなパイプラインの作成と実行を自動化します。このために、Amazon SageMaker と、サーバーレスのワークフローオーケストレーションサービスである AWS Step Functions を利用します。この SDK では、ワークフローのステップを定義しそれらを連結させる 1 つの短い Python スクリプトを使い、ワークフローの作成を実現しています。AWS Step Functions により、基盤となるインフラストラクチャを管理する必要は一切なく、ワークフローのステップを配置することができます。

AWS Step Functions Data Science SDK によれば、提供される各種のプリミティブにより複雑なパイプラインをゼロから構築できます。またそれだけでなく、モデルのトレーニングとデプロイ用のシンプルな TrainingPipeline ワークフローなどを記述した、構築済みの基本的なテンプレートも提供しています。次に示すコードでは、主にトレーニングエスティメーターと Amazon S3 での入出力ロケーションをいくつかのパラメータで指定しながら、パイプラインを定義しています。

import stepfunctions
from stepfunctions.template.pipeline import TrainingPipeline

workflow_execution_role = "<StepFunctions-execution-role-arn>"

pipeline = TrainingPipeline(
    estimator=estimator,
    role=workflow_execution_role,
    inputs=inputs,
    s3_bucket=bucket
)

パイプラインの定義を終えると、それをグラフとして可視化したり、インスタンス化したり、あるいは、反復的に実行できるようになります。実際、複数のワークフローを並列的に実行することも可能です。実行中のワークフローは、AWS Step Functions コンソールを使うか、パイプラインの render_progress メソッドを呼び出すことで、その進行状況をチェックできます。次の図に、トレーニングステップにおいて進行中のワークフローの実行状況を示しています。

これ以外にも The AWS Step Functions Data Science SDK では、TensorFlow 2 や他の機械学習プロジェクトを自動化する、ワークフローが利用可能です。その一例は、モデルの反復的な再トレーニングを自動化するワークフローです。このようなワークフローには、トレーニング後のモデル品質についてのテストが含まれます。これには、品質テスト (モデルはデプロイ済み) 、あるいは失敗の場合 (モデルは未デプロイ) ごとに後に得られた条件が使用されます他のワークフローステップの例としては、自動モデルチューニング、AWS Glue を使う ETL、その他が考えられます。再トレーニングワークフローの詳細については、「Automating model retraining and deployment using the AWS Step Functions Data Science SDK for Amazon SageMaker」をご参照ください。

まとめ

今回の記事では、データ変換、トレーニングと推論コードのプロトタイピング、自動モデルチューニング、およびホスト型トレーニングと推論のための、Amazon SageMaker の機能について解説しました。加えて、AWS Step Functions Data Science SDK を活用して、プロジェクトのプロトタイピングがが完了した後にワークフローを自動化する方法も説明しました。これらの機能は、Amazon SageMaker で TensorFlow 2 や他の深層学習フレームワークを使うプロジェクトにとって、主要な要素となります。

また、こういった機能の他にも、おおくの利用方法があります。たとえば、一過性や展開性の重みと言った、モデルトレーニングにおける一般的な問題に対しても、Amazon SageMaker Debugger は利用価値があります。実稼働にデプロイ済みのモデルにおけるデータドリフトなどの共通的な問題には、Amazon SageMaker Model Monitor も適用できます。Amazon SageMaker ワークフロー機能についての詳細は、関連の GitHub レポジトリをご参照ください。

著者について

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

Amazon Web Services ブログ