Amazon SageMakerで独自アルゴリズムを使って学習してみる

SageMaker について

公式
https://aws.amazon.com/jp/blogs/news/amazon-sagemaker/

ざっくり言うと学習・モデルデプロイ・ホスティングまでできるマネージドサービスです。
組み込みアルゴリズムがありますが、勿論独自アルゴリズムも使えます(要Docker image)

この SageMaker を最近触り始めました。今回は「独自アルゴリズムを使った学習」について
試したことや少し工夫したことを書きたいと思います。

独自アルゴリズムでの学習方法

decision tree を例としたサンプルがあります。これを参考にすれば良いと思います。
https://github.com/awslabs/amazon-sagemaker-examples/blob/master/advanced_functionality/scikit_bring_your_own/scikit_bring_your_own.ipynb

ざっくりまとめると これらの必要ファイル のうち、 train を独自アルゴリズムに改修して
Docker image 内に配置すれば良いようです。

尚、モデルトレーニングの場合 SageMaker は docker run image train を
実行するらしいので、 train ファイルはその名前のままで、実行権限を付与しておく
必要がありそうです。(Dockerfile でエントリーポイント設定すれば話は別)
https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/your-algorithms-training-algo.html

取り敢えずやってみる

以下のようなスクリプトを作成し、トレーニングジョブを実行してみました。
（irisデータ分類の例となります）

sage_test.py

import boto3

client = boto3.client('sagemaker')

training_params = {
    "TrainingJobName": 任意のジョブ名（AWSアカウント内で unique である必要あり）,
    "HyperParameters": {  # 学習(今回は分類)時のパラメータ
            'objective': 'multiclass',
            'num_class': '3'
    },
    "AlgorithmSpecification": {
        'TrainingImage': 学習で使用する Docker image,
        'TrainingInputMode': 'File'
    },
    "RoleArn": SageMakerにアタッチするロール(SageMakerFullAccess があればOK),
    "InputDataConfig": [
        {
            'ChannelName': 'training',  # 任意。train 内でもチャンネル指定するので注意
            'DataSource': {
                'S3DataSource': {
                    'S3DataType': 'S3Prefix',  # s3上のファイルを使う場合はこれ
                    'S3Uri': トレーニングデータのあるs3バケットパス(オブジェクト直指定ではない)
                }
            }
        }
    ],
    "OutputDataConfig": {
        'S3OutputPath': モデルを出力するs3パス
    },
    "ResourceConfig": {
        'InstanceType': 'ml.m4.xlarge',  # インスタンスタイプ
        'InstanceCount': 1,  # 学習インスタンス台数
        'VolumeSizeInGB': 10  # 学習インスタンスのボリューム
    },
    "StoppingCondition": {
        'MaxRuntimeInSeconds': 60*60  # この時間内に終わらなかったらエラー終了(単位:秒)
    }
}

response = client.create_training_job(**training_params)

status = client.describe_training_job(TrainingJobName=jobname)['TrainingJobStatus']
print(status)

try:
    client.get_waiter('training_job_completed_or_stopped').wait(TrainingJobName=jobname)
finally:
    status = client.describe_training_job(TrainingJobName=jobname)['TrainingJobStatus']
    print("Training job ended with status: " + status)
    if status == 'Failed':
        message = client.describe_training_job(TrainingJobName=jobname)['FailureReason']
        print('Training failed with the following error: {}'.format(message))
        raise Exception('Training job failed')

いくつか注意

• ChannelName を設定した場合、SageMakerでは /opt/ml/input/data/[ChannelName] にトレーニングデータが配置されます。
  というか train のこの辺(↓)はおまじないと思って毎度設定しておけば
  良いかもしれないですね。

train内にて

# These are the paths to where SageMaker mounts interesting things in your container.

prefix = '/opt/ml/'

input_path = prefix + 'input/data'
output_path = os.path.join(prefix, 'output')  # 実行結果(成功/失敗)出力先
model_path = os.path.join(prefix, 'model')  # モデル出力先、s3に反映される
param_path = os.path.join(prefix, 'input/config/hyperparameters.json')  # パラメータを設定した場合ここに置かれる

# This algorithm has a single channel of input data called 'training'. Since we run in
# File mode, the input files are copied to the directory specified here.
channel_name='training'
training_path = os.path.join(input_path, channel_name)  # トレーニングデータはここに置かれる

• HyperParameters は渡す際は key, value ともに str 型 の dict にする必要があります。
  なので、もしも学習時に他の型でパラメータを使いたい場合は変換する必要があります。

train内にて

        # Here we only support a single hyperparameter. Note that hyperparameters are always passed in as
        # strings, so we need to do any necessary conversions.
        max_leaf_nodes = trainingParams.get('max_leaf_nodes', None)
        if max_leaf_nodes is not None:
            max_leaf_nodes = int(max_leaf_nodes)  # int 型に変換

自分なりの改善

この手順通りにやれば割と簡単に SageMaker で学習ができました。
一方で、私の中では

このサンプル通りに train にアルゴリズムを書く方法をとっていたら
その度に Docker image を作る必要があり、image が乱立して管理が
面倒なことになりそうだな…

とも思いました。

そこで、この問題解決のために以下のように変更を加えてみました。

• Docker 内の構成を以下のように変更
  (train_prg フォルダを作り、その配下に独自アルゴリズムのpythonスクリプトを配置)

/opt/program/
    ├── nginx.conf
    ├── predictor.py
    ├── serve
    ├── train
    ├── wsgi.py
    └── train_prg/　　# このフォルダを追加
　　　　　　　　　　　　　　　　　　　└── training_1.py

• sage_test.py の HyperParameters に実行スクリプト名を設定

sage_test.pyにて

    "HyperParameters": {
            'prg': 'training_1.py',  # 実行スクリプト名を追加
            'objective': 'multiclass',
            'num_class': '3'
    },

• train 内で目当ての pythonスクリプトを実行

train

import json
import subprocess

with open(param_path, 'r') as tc:
    prg = json.load(tc)['prg']

subprocess.run(['python', 'train_prg/{}'.format(prg)])

• train_prg/training_1.py は 旧 train と同じ内容とする(独自アルゴリズム本体)

まとめるとこういうことをしています。

• train は python スクリプト呼び出し(親スクリプト)の役割に専念させた
• 実行する独自アルゴリズム(pythonスクリプト)は train_prg 配下に配置した
• 独自アルゴリズムが増えた場合は train_prg 配下に配置する
  → スクリプトファイルなんて沢山置いてもそう容量増えないだろうという想定
• 実行する独自アルゴリズムは HyperParameters で制御する
• Docker image には全アルゴリズムの実行に必要なパッケージを入れる

こうすることにより、SageMaker に必要な Docker image は 1 つだけとし、乱立を防ぐことが出来そうです。(サイズが大きくなる可能性はある)

感想

日本語のドキュメントが少なく、この程度の内容でもここまで理解するのが大変でした。
やっぱり英語できたほうが良いなぁ…

理解を誤っているところもあると思うので、見つけ次第改善していきたいと思います。

今はまだ学習部分しか出来ていませんが、推論部分もいずれ書きたいと思います。