【Techの道も一歩から】第41回「SageMakerでStep Functionsを使った学習パイプラインを構築するために概念を理解する」

こんにちは。 DSOC R&D グループの高橋寛治です。

最近は、MLOpsへの興味が大きく、ツールの実利用や勉強会での発表を行っています。

機械学習プロジェクトにおいて、久々の実行や引き継いだコードの実行の際に、順序や実行環境を思い出したり準備したりしながら確認するのはすごく大変です。コンテナ化された上でパイプラインがあると、実行順序も実行環境も覚えたり記録したりする必要がなくなります。

コンテナとパイプラインを組み合わせて利用するAmazon SageMakerを利用し、学習パイプラインを構築してみました。その概要とAmazon SageMakerを利用したパイプラインを構築する上で押さえておくべきだと思う概念を簡単に紹介します。

SageMakerによるパイプライン構築の概要

AWS Step Functions Data Science SDK for Amazon SageMakerを利用して、Step Functions上で学習パイプラインを構築します。

Step Functionsとは、マネージドのワークフローサービスです。このワークフロー内の一つの作業単位をタスクと呼び、タスクは一つのAWSサービスが実行されます。例えば、Athenaでの集計というタスクや、ECSによるバッチ処理というタスクとなります。 Step Functionsのステートマシン（ワークフロー）は、ASL（Amazon States Language）というJSONベースの言語で定義されます。

AWS Step Fucntions Data Science SDK for Amazon SageMakerは、SageMakerとStep Functionsを組み合わせて機械学習モデルの学習やデプロイといった機械学習関連のワークフロー構築を簡単にするライブラリです。 Step Functionsのステートマシンを書くためにASLを使う必要が無く、親しみのあるPythonコードとして記述することができます。

作成したパイプライン

たびたびブログのテーマにも上げている固有表現抽出を題材としました。

以下は、Step Functions上で可視化されたシンプルなワークフローです。

f:id:kanjirz50:20210823204341p:plain — シンプルなBERT Training Pipeline

ワークフローの流れを簡単に紹介します。

前準備として、あらかじめS3に学習・評価・開発用のアノテーション済みデータをPUTしておきます。

ステートマシンの実行を開始すると、まずはタスク「SageMaker BERT Preprocessing step」が起動します。ここでは、データに対する前処理を行います。具体的には、トークン化や入力トークン長の制限などを行います。

次に、「SageMaker BERT Training step」が起動します。ここでは、前段のタスクで前処理されたデータを用いてモデルの学習を行います。学習されたモデルはS3にPUTされます。

以上が構築したワークフローです。

AWS Step Fucntions Data Science SDK for Amazon SageMakerを使ってタスクにSageMakerを利用した場合は、その実行の詳細が次のように記録されます。一枚目のスクリーンショットは前処理、二枚目のスクリーンショットは学習の記録を表します。

f:id:kanjirz50:20210823205333p:plain — SageMaker Processing

f:id:kanjirz50:20210823205346p:plain — SageMaker Training

このように、どのようなコンテナでどのロールを使い、どのインスタンスでどれくらいの実行時間がかかったか、といったことがすべて記録として残ります。もちろん入出力に利用したデータのS3のURIやコンテナのエントリポイントおよび引数も記録として残ります。

入出力のデータは、コンテナ起動時にコンテナ内の /opt/ml/data のような指定したディレクトリにS3から自動的に取得されます。つまりコードを書く際にS3であることを意識することなく、ローカルストレージのデータを取り扱う記述をすることになります。

カプセル化された概念を押さえる

SageMakerやStep Functions、またこれらをPythonで取り扱うことのできるAWS Step Functions Data Science SDK for Amazon SageMakerを理解して使う上で、押さえておくべきだと感じた点を紹介します。

まずは、Step FunctionsとStep Functions上で扱えるSageMakerで、何がカプセル化・抽象化されているのかを押さえます。

Step Functionsはワークフローを管理します。そしてワークフローの各ステップ（タスク）はAWSサービスが実行されます。ワークフローには各タスクにどのような入出力を渡すかも示されます。

そのワークフローを構成するタスクとしてSageMakerを利用します。 SageMaker Python SDKのドキュメントの冒頭で何がカプセル化されているかまとまっています。 AWS Step Functions Data Science SDK for Amazon SageMakerではステップとして提供されており、SageMakerの各機能を実行できます。

ProcessingStep：SageMaker Processing Jobを実行するタスク
TrainingStep：SageMaker Training Jobを実行するタスク
TransformStep：SageMaker Transform Jobを実行するタスク
TuningStep：SageMaker HyperParameterTuning Jobを実行するタスク
ModelStep：SageMakerのモデルを作成するタスク
EndpointConfigStep：SageMakerのエンドポイント設定を作成するタスク
EndpointStep：SageMakerのエンドポイントを作成または更新するタスク

そして、これらの各StepのうちJobを実行するものについては次の要素を持ちます。

コンテナ
コード
入出力の定義
（引数やハイパパラメータ）

そして、各種データやコードはS3に、コンテナはECRに登録されます。

まとめると、ワークフローは処理順やどのステップを使用するか何を入出力とするかを定義し、各ステップで実際の処理を行います。

SageMakerのProcessingStepとTrainingStep

SageMakerのProcessingStepは前処理や後処理、評価のような処理を行います。 TrainingStepは文字通り学習を行います。

ここで、各ステップのお作法をいくつか押さえておくと良いと思います。

ProcessingStepでは、コンテナのエントリポイントとしてスクリプトを指定して実行します。具体的には、["python3", "opt/ml/processing/input/code/preprocessing.py"] のように指定します。実行環境には、任意のコンテナを指定します。例えば、前処理のためにライブラリが必要という場合に、ライブラリがインストールされたコンテナを用意しておきます。自然言語処理だとトークナイズのために任意のライブラリが必要になることが多いです。

TrainingStep向けには、学習と推論の両方に対応するコンテナを用意します。推論部については、以前に書いたこちらが参考になります。

buildersbox.corp-sansan.com

簡単に説明すると、学習と推論でエントリポイントが異なるイメージです。学習時はtrain、推論時にはserve がエントリポイントとして基本的には設定されます。すなわち、各エントリポイントで実行したいことを記述したコンテナを用意します。

serveを書かなくても学習ができるかと思いますが、SageMakerの仕組みをフルに使う場合は serve を書いておくと簡単に推論エンドポイントを構築することができます。 SageMakerモデルでは、学習と推論のコンテナおよび学習済み重みファイルは同じものを利用するためです。

良いモデルを作ることに集中したい

実行順や実行環境の用意というのは大事な要素でありますが、仕組み化できる箇所であるため人力でやる必要はないと思います。良いモデルを作るために前処理を工夫したり、様々なモデルを試して試行錯誤するための、考えを練ることに時間をかけるべきでしょう*1。学習パイプラインツールはそういった本質的なことに取り組むための仕組みであると思います。