Amazon SageMakerのScikit-learnコンテナを使ってみる #reinvent

こんにちは、小澤です。

Amazon SageMaker(以下SageMaker)の標準で用意しているコンテナにScikit-learnが加わっています。 re:Inventで大々的に発表された機能ではないですが、リリースされたのが大体そのタイミングで、 私がre:Invent参加時に現地知ったので、個人的にre:Inventネタ感があるということにしていますw

早速使ってみる

では、早速使ってみましょう。 使い方は基本的に、TensorflowやMXNetなどを利用する際と同様で、以下の2つを用意します。

  • コンテナで処理される学習用やエンドポイント用のスクリプト
  • データを渡して学習や推論を行うノートブック

ええ、全く一緒ですね。 では、記述するスクリプトを見てみましょう。

今回は、以前以下のブログでも使った数字の認識処理をやってみたいと思います。

コンテナ上で処理される学習スクリプト

Scikit-learnでの学習処理を行うスクリプトを書くのは非常に簡単です。 用意するものを以下の2つのみになります。

  • 学習の処理を行うmain部分
  • 推論用にモデルを読み込むmodel_fn関数

では順にみていきましょう。

学習処理用のmain部分

学習処理はPythonスクリプトでは見慣れたmainの処理として記述します。

if __name__ == '__main__':
    # ここに処理を記述

この中で行う処理としては、以下のような主に以下のようなものとなります。

  • 実行時に渡されたパラメータをコマンドライン引数として受け取る
  • 学習に利用するデータを読み込む
  • 学習処理を実行する
  • モデルを保存する

この中でちょっと特殊なのは最初に行う"実行時に渡されたパラメータをコマンドライン引数として受け取る"の部分のみとなります。 その処理内容としては、以下のようになります。

parser = argparse.ArgumentParser()

# ハイパーパラメータとして与える値を引数で受け取る
parser.add_argument('--svm_c', type=float, default=1)

# SageMakerの動作に関する引数
# モデルの出力先パスを指定
parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
# 学習データの入力パスを指定
parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAIN'])

args = parser.parse_args()

Pythonを普段から書いている人にはおなじみのargparseを使って引数を処理しています。 コマンドライン引数をどのように渡すかはこの後ノートブック側の処理で解説しますが、 主にハイパーパラメータとして渡す任意の引数と、SageMaker固有の引数に分けられます。

上記スクリプトでは、--svm_cとしているパラメータがハイパーパラメータとして渡す引数になっており、 オプション名や個数は任意の数決めることが可能です。

SageMaker固有の引数はオプション名も固定で決まっており、以下のものを指定します。

オプション 内容
--model-dir 学習後のモデルの保存先
--output-data-dir モデル以外の出力ファイルの保存先
--train 学習データのパス
--test 学習時にテストチャネルとして指定する項目

これらは不要であれば省略してかまいません。 ほぼ必須となるのは、モデル出力先と学習データのパスでしょうか。

また、上記スクリプトでdefaultとして指定しますが、学習コンテナではこれらのデフォルト値となるような値があらかじめ環境変数で設定されていますので、それを活用するのもいいでしょう。

それぞれのオプションに対する環境変数の対応関係は以下のようになります。

オプション 環境変数
--model-dir SM_MODEL_DIR
--output-data-dir SM_OUTPUT_DATA_DIR
--train SM_CHANNEL_TRAIN
--test SM_CHANNEL_TEST

これらの環境変数が具体的にsagemaker-containersのREADME.mdに記載されています。

コマンドライン引数の取得後は、SageMaker固有のライブラリなどをほぼ使うことなく、通常のScikit-learnの処理として実装可能です。

学習に利用するデータはS3から取得することになるので、以下のように、引数で受け取ったパスから読み込んでいます。

input_files = [ os.path.join(args.train, file) for file in os.listdir(args.train) ]
if len(input_files) == 0:
    raise ValueError(('There are no files in {}.\n' +
                      'This usually indicates that the channel ({}) was incorrectly specified,\n' +
                      'the data specification in S3 was incorrectly specified or the role specified\n' +
                      'does not have permission to access the data.').format(args.train, "train"))
raw_data = [ pd.read_csv(file, header=None, engine="python") for file in input_files ]
train_data = pd.concat(raw_data)

この処理では引数としてディレクトリを受け取り、その中にあるすべてのファイルを対象とする想定で記述されています。

内容を見ていただくとわかる通り、実装方法としてその限りでなくても動きますが、公式のexampleなどもこのような書き方になっていることから統一しておくことをオススメします。

Sckit-learnでは学習時にデータとラベルをそれぞれ別の引数として渡す必要があるので、分離する処理も行っています。

# 読み込んだデータを学習データと正解ラベルに分ける
train_y = train_data.ix[:,0]
train_x = train_data.ix[:,1:]

これで学習ができる状態になりました。 学習をする際の処理も通常のScikit-learnの使い方となります。

clf = SVC(C=args.svm_c)
clf.fit(train_x, train_y)

ハイパーパラメータは引数で受け取った値を利用しています。

学習が完了したのち、推論エンドポイントで利用するために、モデルをファイルとして保存しておく必要がります。 こちらも通常のScikit-learnと同様、joblibを使って保存してやるのみです。

joblib.dump(clf, os.path.join(args.model_dir, "model.joblib"))

これで学習処理を実行するための処理の記述は完了となります。

推論エンドポイント用の関数を作成する

最後に推論エンドポイントから利用できるようにモデルを読み込む処理を記述した関数を用意しておきます。

用意する関数はmodel_fnとなります。 特別な処理が必要なければ、ここは先ほど保存したモデルを読み込むのみになります。

def model_fn(model_dir):
    """
    Predictする際に利用するためにモデルを読み込む
    """
    clf = joblib.load(os.path.join(model_dir, "model.joblib"))
    return clf

データを渡して学習や推論を行うノートブック

他のフレームワークが入ったコンテナを使った学習同様、こちらも先ほどのスクリプトを使って実際に学習する処理を記述する必要があります。 こちらの処理もほぼ同じような記述となります。

まず、SageMakerのセッションの取得と、S3の保存先を定義しておきます。

sagemaker_session = sagemaker.Session()
role = get_execution_role()

s3_bucket = '<Your S3 Bucket>'
prefix = 'sklearn-test/train'

利用するデータを取得してS3に保存しておきましょう。

digits = load_digits(10)
# データを学習用とテスト用に分ける
train_x, test_x, train_y, test_y = train_test_split(digits['data'], digits['target'], test_size=0.3, random_state=71)

# 正解ラベルも含めたデータになっている必要があるため、1列目に正解ラベルを持つ形式にする
# どのような形式で出力する必要があるかは先ほど記述した学習処理の際の読み込み方に依存する
joined_data = np.insert(train_x, 0, train_y, axis=1)

os.makedirs('./data', exist_ok=True)
np.savetxt('./data/train_digits.csv', joined_data, delimiter=',')
train_input = sagemaker_session.upload_data('./data/train_digits.csv', bucket=s3_bucket, key_prefix=prefix)

データの準備ができたので、Scikit-learnを使うためのEstimatorインスタンスを作成します。

from sagemaker.sklearn.estimator import SKLearn

sklearn = SKLearn(
    entry_point='sklearn_test.py',
    train_instance_type="ml.c4.xlarge",
    framework_version='0.20.0',
    role=role,
    sagemaker_session=sagemaker_session,
    hyperparameters={
        'svm_c' : 1
    },
    output_path='s3://<Your S3 Bueckt>/sklearn-test/model'
)

entry_point引数で先ほど作成したスクリプトを指定しています。 また、スクリプト内でコマンドライン引数として受け取るもの二関してはここで設定しています。

hyperparameters引数でdict形式で任意の引数となる項目の値を指定します。 ここでは--svm_cオプションに対してsvm_cのように"--"を除いた形式で指定します。

また、モデルやファイルの出力先のパスをデフォルトから変更するためにoutput_path引数でS3のパスを指定しています。 今回は、--output-data-dirオプションに渡される値は利用していないので、--model-dirのデフォルト値がここで指定した値の下にジョブ名の子階層が切られてモデルが出力されます。

あとはfit関数を呼び出せば学習の処理が実行されます。

sklearn.fit({'train': train_input})

スクリプトの--trainおよび--testオプションに渡す値はここで指定します。

学習が完了するとあとはdeploy関数でエンドポイントを作成して、predict関数で未知のデータに対する推論が可能となります。

predictor = sklearn.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge")
predictor.predict(test_x)

動作確認後、不要な場合はエンドポイント削除を忘れないようにしましょう。

sklearn.delete_endpoint()

おわりに

今回は、SageMakerでScikit-learnを使った機械学習をする方法を解説しました。

実際使ってみると、ほぼほぼいつものScikit-learnを使う感じでそのまま利用できるので非常に便利です。