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

この間、社内で、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`とするやり方が多かったが、
自分的に、そのやり方が気持ち悪かったので、

```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側で、ノートブックインスタンスを立ち上げて、学習を掛けてみた。
学習をするコードは、以下のような感じ。

```.py
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を通してから実行するようにした。

```train
#! /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`を変更して実行してみた。

```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/`してから、学習ソースの実行を行うようにするといいかもしれない。