SageMakerカスタムコンテナでの学習ソースの外挿 #Docker

この間、社内で、SageMakerを使って行こうという話になったので、自分が今やっているものをSageMakerへ載せ替えしようとした。
ところが、自分はKerasを使っていて、デフォルトではサポートされていなかった。
そこで、どの様にやるのか調べたところ、デフォルトでサポートされたフレームワーク/アルゴリズムの記事は、いっぱい出てきたが(ほとんどハンズオンの書き写しだったが...)
カスタムコンテナの情報になると、なかなか記事が見つからず困った。
公式ドキュメントもよくわからない感じだった（おそらく自分の読解力のなさによる)
いくつか引っかかりながらも、学習まではできるようになったので、そのメモ。

また、
カスタムコンテナを使用する場合、Dockerコンテナ内に学習ソースを入れていないとできないようで、
ネットワークの修正等をしようとする度に、別でDocker立ち上げて、中身変えて、commit, pushしないといけないようで、
ネットワークをよくいじる自分からすると非常に面倒だった。
そこで、カスタムコンテナでも、一度作ったDockerを変更しなくても、外部から学習ソースを渡す(外挿する)方法を見つけたので、そのメモ。

カスタムコンテナの用意

今回作ったDockerコンテナは、
Ubuntuをベースにして、pyenv　+ python3.6を入れたものを用意した。
内部のディレクトリ/ファイル構成は、このような感じ。

/opt/program/train
    /ml/output
       /model
       /input/data/train_inputs
                  /train_teachers
                  /valid_inputs
                  /valid_teachers
                  /src/train.py

ちなみに、WORKDIRを/opt/program に指定している。

ここで、
SageMakerは、学習開始時に、

$ docker run <イメージ名> train

として学習を実行する。
よって、学習ソースを
WORKDIR以下のtrainというファイルにしておく。

ただ、pythonソースから.pyを取って、trainとするやり方が多かったが、
自分的に、そのやり方が気持ち悪かったので、

python src/train.py

を実行可能形式で用意し、実際のソースを
/opt/program/src/train.py
に用意した。

また、
/opt/ml/input/data以下は、SageMakerノートブックインスタンスでS3からデータ転送を行える領域になる。
input以下のディレクトリ名自由につけられ、S3転送時に転送先ディレクトリとしてキーワードで指定すれば、そこにデータが転送される。

/opt/ml/modelは、ここ以下に学習済みモデルを置くことで、S3側に転送される。

このようなカスタムコンテナを用意し、docker commitの後に、ECSにdocker pushしておく。
ECSへのコンテナのpush方法は、ECSでリポジトリを作ると、そこで確認できる。

ノートブックインスタンスからの学習起動コード

カスタムコンテナが用意できたら、SageMaker側で、ノートブックインスタンスを立ち上げて、学習を掛けてみた。
学習をするコードは、以下のような感じ。

import os
from datetime import datetime as dtime
from sagemaker import get_execution_role
import sagemaker as sm

PROJECT_NAME = 'sagemaker-keras'

S3_BUCKET = 's3://box-side-segmentation'
INPUT_BUCKET = S3_BUCKET + '/input'
OUTPUT_BUCKET = S3_BUCKET + '/output'
MODEL_BUCKET = S3_BUCKET + '/model'

KRY_TRAIN_INPUT = 'train_inputs'
KRY_TRAIN_TEACHER = 'train_teachers'
KRY_VALID_INPUT = 'valid_inputs'
KRY_VALID_TEACHER = 'valid_teachers'

INSTANCE_COUNT = 1
INSTANCE_TYPE = 'ml.p2.xlarge'

def get_timestamp():
    n = dtime.now()
    return '%04d-%02d%02d--%02d%02d-%02d' % (n.year, n.month, n.day, n.hour, n.minute, n.second)

role = get_execution_role()
sess = sage.Session()
account = sess.boto_session.client('sts').get_caller_identity()['Account']
region = sess.boto_session.region_name
image = f'{account}.dkr.ecr.{region}.amazonaws.com/{PROJECT_NAME}:latest'

inputs = { KRY_TRAIN_INPUT : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_TRAIN_INPUT + '/')
         , KRY_TRAIN_TEACHER : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_TRAIN_TEACHER + '/')
         , KRY_VALID_INPUT : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_VALID_INPUT + '/')
         , KRY_VALID_TEACHER : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_VALID_TEACHER + '/')
         }
output_path = S3_BUCKET + '/output'

print('image :', image)
print('inputs : ', inputs)
print('output : ', output_path)

timestamp = get_timestamp()

classifier = sm.estimator.Estimator(image, role, INSTANCE_COUNT, INSTANCE_TYPE
                                   , output_path=output_path
                                   , sagemaker_session=sess)

classifier.fit(inputs, job_name='box-side-seg-sample' + timestamp)

　
学習用データのS3から学習用docker内への転送の指定は、以下の箇所で行っている。

inputs = { KRY_TRAIN_INPUT : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_TRAIN_INPUT + '/')
         , KRY_TRAIN_TEACHER : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_TRAIN_TEACHER + '/')
         , KRY_VALID_INPUT : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_VALID_INPUT + '/')
         , KRY_VALID_TEACHER : sm.session.s3_input(INPUT_BUCKET + '/' + KRY_VALID_TEACHER + '/')
         }

上記のように、
inputsに、キーを転送先ディレクトリ名にして、値として転送元(sm.session.s3_inputで指定)をもつ辞書を指定することで、データの転送ができる。

python2.7が動く問題

と、実行してみたところ、なぜか、学習がSyntax Errorを吐いて落ちた。
よく見たらpython2.7で実行されている模様。
なぜだ？と思ったら、docker実行時に/bin/bashではなく、直接trainが実行され、pythonが動くので、pyenvの環境設定が聞いてない様子。

そこで、以下のようにtrain内の最初に、環境設定をしてPathを通してから実行するようにした。

#! /bin/sh
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"

pyenv rehash
python -V
python src/train.py

これで、python3.6 で、実行され、学習がされることの確認ができた。

学習ソースの外挿

ここからが、+αの内容。
学習ソースを変更するたびに、Dockerを立ち上げて、変更、commit、pushを行うのは非常に面倒な気がした。

そこで、ふと思ったのが、
- 実行されるtrainの中に書いてある任意のコマンドが実行可能なので、好きな位置のpythonファイルの実行が可能。
- 学習前に、inputsに指定することで、(/opt/ml/input/data以下ではあるが)任意のファイルを転送可能。
上の２点を合わせれば、
S3上に学習ソースをおいて、学習時に転送することで、外挿できるのでは？

そこで、以下のようにtrainを変更して実行してみた。

#! /bin/sh
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"

pyenv rehash
python -V
python /opt/ml/input/data/src/train.py

これで、思惑通り、カスタムコンテナでも、学習ソースを外挿可能になり、
変更のたびに面倒なDockerの更新が不要になった。

[追記]
train内では、任意のlinuxコマンドが実行可能なので、
特に、学習ソースを複数のファイルに分けている場合、importの関係上、
cd /opt/ml/input/data/src/してから、学習ソースの実行を行うようにするといいかもしれない。