SageMakerの汎用的なdockerイメージ #AWS

はじめに

こんにちは。
かなり遅れましたがSageMakerネタ二回目です。

SageMakerへのトレーニングジョブ発行する方法はアルゴリズムを指定するだけの簡単なものからフレームワークを利用するもの、そもそもアルゴリズムを自由に記述する方法があります。
今回はその中でも自由なアルゴリズムをどう記述するのかをみていきます。

トレーニングジョブの仕組み

詳細までみていくと割とややこしいので概要から少しずつ深堀していきます。

概要

基本的な機械学習をするにあたって必要になるものは

トレーニングデータなどのデータ
モデルの学習アルゴリズムのソース
計算リソース

かなと思います、それをSageMakerでは

自分で用意してs3に配置
学習アルゴリズムのソースが含まれるdockerイメージ
計算インスタンス

このように扱います。
学習手順としては(詳細は異なるかもしれませんが)以下のようになります。

指定タイプのインスタンスを起動
指定データをs3からダウンロード
指定dockerイメージから、データやその他ファイルをマウントしイメージ内の学習スクリプトをコンテナとして起動
学習結果(モデル)をs3にアップロード

基本的にこれしかやってません。簡単ですね。

docker関連はどうなってんの？

一番ブラックボックスなのはここだと思います。
ここの仕様を把握しないとお目当て(タイトル)のものは作れません。
そこでトレーニングジョブでどのようにイメージが使われているのか調べました。

起動コマンド

まずはSageMakerがどうやってイメージを起動しているのか知る必要があります。
ここにこのように書かれています。

docker run image train

これで何かしらtrainというものが実行されるか、あらかじめイメージ内で指定したENTRYPOINTの引数としてtrainが渡されることがわかります。
「とりあえずtrainって名前のスクリプト書いてパスの通ったところに置いとけばいいんだな」と認識すればなんとかなります。

コンテナのディレクトリ構成

次に気になるのが、データの「ありか」です。
これはSageMaker側が勝手にs3からダウンロードしてマウントするのでこれについてもどこにマウントしているのか仕様を把握する必要があります。

調べました。

/opt/ml/  
     ├ input/  
     │   ├ data/  
     │   │   ├ CHANNEL_NAME/
     │   │   │    ︙
     │   │   └ CHANNEL_NAME/
     │   └ config/ 
     │      ├ hyperparameters.json  
     │      ├ inputdataconfig.json  
     │      └ resourceconfig.json  
     ├ output/  
     │   └ failure  
     └ model/

簡単にそれぞれの役割を説明します。

data/CHANNEL_NAME/: チャネル(下で補足有)で指定したディレクトリ名以下に指定したs3パス以下のファイルが置かれる(CHANNEL_NAMEは指定したチャネル名に置き換わる)
hyperparameters.json: 指定したハイパーパラメータ(後述)がjson形式で置かれる
inputdataconfig.json: 入力データのchannelの情報などが記されたjsonが置かれる
resourceconfig.json: クラスタの構成．ホスト名などが書かれたjson。分散処理用
output/: ここに置いたものはs3のoutput先にoutput.tar.gzとして圧縮されてアップロードされる
failure: 処理が失敗したときにここに処理を書くといいらしい
model/: ここに出力されたものはmodel.tar.gzととしてs3のoutput先にアップロードされる

これでtrainって名前のスクリプトにだいたい何を書けばいいのか決まってきました。
「インプットされたデータを加工して、結果を出力用のディレクトリに置けばいいんだな」ということです。

チャネルについての補足

これはSageMaker特有の概念的なあれだとおもいます。
要するに複数箇所の入力データ用のs3をマウントするときにそれぞれ名前をつけて区別しようというやつです。

具体的な利用方法としては予めデータをトレーニング用と検証用に分けて

s3://bucket/.../train/
s3://bucket/.../validation/

のように配置しchannel名に「train」・「validation」とするとか

他にはある程度学習したモデルと、追加データを用意して

s3://bucket/.../model/
s3://bucket/.../train/

のように配置しchannel名に「model」・「train」とし転移学習・増分学習するとか。

使い方は自由です。いい感じに指定してください。
ただし、指定できる数は８個までのようです。

ハイパーパラメータ

入力できるのはs3のデータだけではありません。
ハイパーパラメータという~~名目で~~任意の文字列をコンテナに渡せます。逆にこれら以外にコンテナに干渉する方法がないです。実態は先ほどディレクトリ構成で説明したように、キー・バリューがjson形式で置かれるだけですが…

これをコンテナ内でうまく処理することで実質環境変数のように使ったりといった用途が考えられます。もちろんアルゴリズムの正しい意味でのハイパーパラメータとしても使います。

まとめ

トレーニングジョブ用のイメージは以下を満たせば最低限動きます。

trainというコマンドもしくは引数を受け取れるようにする
/opt/ml/modelにモデルファイルを出力する

API提供の仕組み

ここはさっくりと説明します。

概要

予測のAPIサーバの提供もdockerイメージで行われています。
トレーニングとは別のイメージを使うこともできるようですが、同じものを使うものが多そうです。
フローイメージとしては以下になります。

指定タイプのインスタンスを起動
指定モデルデータをs3からダウンロード
指定dockerイメージから、モデルデータやその他ファイルをマウントしイメージ内のサーバをコンテナとして起動

docker関連

こちらに要件は書かれています。

まとめると以下のようになります。

docker run image serveで起動されう
/opt/ml/modelディレクトリからモデルを取得できる
API要件
- ポート番号8080で/pingでGETを受け取りHTTP200を返す
- ポート番号8080で/invocationsにPOSTを受け取り推論結果を返す
ただし、学習と違ってハイパーパラメータで情報を渡すのではなくこちらでは環境変数を使う

任意のアルゴリズムを受け付ける

とりあえずdockerイメージの要件はわかったので入力と出力さえ気をつければ任意のアルゴリズムのイメージは作成できると思います。
これで終わりでもいいのですがもう少し踏み込んで、あらかじめ作ったdockerイメージを変えず任意のアルゴリズムを流し込む方法を考えます。

tensorflowやchainerなどのフレームワークを扱える公式のイメージではentrypointと呼ばれるスクリプトファイルを後付けで渡すことができます。
今回はこれを参考にします。

どうやっているのか

フレームワーク系の公式イメージはどうやってコンテナを立てた後にローカルのスクリプトファイルを受け取っているのでしょうか？

答えは単純でs3にアップした後にコンテナ内でダウンロードをしているにすぎません。
しかし、コンテナはs3にアップしたファイルの場所を知らないため別の方法で教えてあげる必要があります。
ここで、ハイパーパラメータや環境変数を利用しているわけです。
特別な名前のハイパーパラメータや環境変数名をダウンロードすべきs3パスと取り決めてやれば可能となります。

フレームワーク系のトレーニングジョブを発行し詳細をGUIから確認してみると
sagemaker_xxxxxという風に自分で設定したハイパーパラメータ以外のものが渡されていることがわかります。
sagemaker_submit_directoryはまさに自前のentrypointを含むディレクトリがアップされたs3パスになっています。
公式のフレームワークイメージではsagemaker-containersのパッケージを使ってこれらのハイパーパラメータを読み取ったりs3からスクリプトをダウンロードしたりを行っています。

汎用的に

この仕組みを真似れば任意のスクリプトをコンテナに後付けで実行させられます。

スクリプトをアップロードしたs3のパスと起動スクリプト名をハイパーパラメータまたは環境変数として渡す
s3からファイルをダウンロードして、その中の起動スクリプトを実行

docker run image trainが実行されればハイパーパラメータを読み取り、docker run image serveが実行されれば環境変数を読み取れば一つのイメージで完結できます。

イメージを使う

作ったイメージの使い方は簡単です。

作ったイメージをECRにアップロード
アルゴリズムを書いたスクリプトをs3にアップロード
SageMakerのトレーニングジョブの作成のGUIからECR での「独自のアルゴリズムコンテナ」を選び先ほどのイメージの場所を指定
必要な情報をハイパーパラメータにセット
あとは普通のトレーニングジョブと同じです

もちろん、これらはsagemakerのsdkやboto3を使ってもできます。

ただ、ここでs3にスクリプトをアップロードしたり、そのパスをハイパーパラメータにセットしたりが面倒なので、この辺を自動化するように工夫するのもいいかとおもいます。
sdkの公式のフレームワーク用のEstimatorクラスでは自動アップロード・自動ハイパーパラメータセットするようになっているのでそれを参考にするのがいいと思います。

最後に

ざっくりとSageMakerのイメージがどのように実行されるかを紹介しました。基本的に公式イメージで事足りると思いますがdockerの勉強にもなるので是非オリジナルのイメージ(とそれを楽に使う仕組み)を作って遊んでみてください。