SageMakerのトレーニングジョブでマネージドスポットトレーニング機能がリリースされました！

osawa yuto

2019.08.27

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

どうも、DA事業本部の大澤です。

SageMakerでのモデルのトレーニングジョブにおいてマネージドスポットトレーニング機能がリリースされました!

モデルのトレーニングジョブでスポットインスタンスを活用することができ、トレーニングジョブにおけるインスタンス費用を最大90%節約できます!

Managed Spot Training: Save Up to 90% On Your Amazon SageMaker Training Jobs | AWS News Blog

※ マネージドスポットトレーニングはスポットインスタンスを活用する機能となるため、ジョブの遅延や中断、停止の可能性があります。使用用途/環境においてそういったリスクを許容できるかどうかの十分な検討が必要です。

概要

Amazon SageMaker

Amazon SageMaker は、すべての開発者とデータサイエンティストに機械学習モデルの構築、トレーニング、デプロイ手段を提供します。Amazon SageMaker は、機械学習のワークフロー全体をカバーする完全マネージド型サービスです。データをラベル付けして準備し、アルゴリズムを選択して、モデルのトレーニングを行い、デプロイのための調整と最適化を行い、予測を行い、実行します。モデルをより少ない労力と費用で、本番稼働させることができます。

Amazon SageMaker（機械学習モデルを大規模に構築、トレーニング、デプロイ）| AWSより

SageMakerは機械学習のフルマネージドサービスです。モデルの学習やデプロイ、ノートブックのホストなどがインフラの管理なく、可能となります。

今回はSageMakerのモデルのトレーニングジョブにおける機能追加となります。

マネージドスポットトレーニング

マネージドスポットトレーニングはモデルのトレーニングジョブでスポットインスタンスを利用できる機能です。

SageMakerはフルマネージドサービスなので、スポットインスタンスを起動して...みたいなことをする必要はありません。マネージドスポットトレーニングを有効化することで自動的にトレーニングジョブで使用するスポットインスタンスがプロビジョニングされます。

マネージドスポットトレーニングを有効にするには、トレーニングジョブ作成時にマネージドスポットトレーニングを有効にし、Maximum wait time(トレーニングジョブが完了するまでにどれだけ待ち続けるかの期間)を指定するだけです。

後述するジョブの中断に備えて、チェックポイント用のS3URIとチェックポイント用のローカルパス)を指定することもできます。

インスタンス費用の節約率

マネージドスポットトレーニングにおけるインスタンス費用は(Billable Time * インスタンスの時間単価)で計算できます。実際の学習時間であるTraining TimeとBillable Timeの差とTraining Timeの比が、マネージドスポットトレーニングにおけるインスタンス費用の節約率となります。

すなわち、(1 - Billable Time / Training Time) * 100が節約率となります。

例

(1 - 70 / 279) * 100 ≒ 75%

...といった感じでインスタンス費用の節約率が計算できます。

トレーニングジョブの中断とジョブが長引く可能性

マネージドスポットトレーニングはスポットインスタンスを活用するため、インスタンスの空き状況によってはトレーニングジョブの開始が遅れたり、ジョブが中断し、結果的にジョブの実行時間が長引く、もしくは途中で終了する場合があります。トレーニングジョブが完了せず終了したり、ジョブの実行が長引くなどのリスクが許容できない場合には、マネージドスポットトレーニングの有効化は避けておいたほうが良さそうです。

トレーニングジョブが中断した場合でも待機/中断時間を含めたジョブの実行時間がMaximum wait time以内かつ、チェックポイント用のS3URIが指定してあれば、インスタンスが確保されるまでジョブは待機します。Maximum wait timeを超えたり、チェックポイント用のS3URIが指定してなければ、ジョブは中断した時点で終了します。

特段の理由がなければ、チェックポイント用のS3URIは設定しておいたほうがよさそうです。

※ 2019年8月27日現在、組み込みアルゴリズムのうち、チェックポイントに対応しているのは画像分類、物体検出、セマンティックセグメンテーションのみです。それ以外の組み込みアルゴリズムとチェックポイントを行わないマーケットプレイスのアルゴリズムではMaximum wait timeは3600秒を上限とする制限があります。

やってみた

早速、マネージドスポットトレーニングを試してみます。先ほどマネジメントコンソールでの設定画面を紹介したので、ここではSageMaker Python SDKを使って進めたいと思います。

MNISTのデータセットを使って、手書き数字画像が0かそれ以外かを分類する二値分類モデルを作成します。学習までの手順は基本的に以下のノートブックに則った内容です。

amazon-sagemaker-examples/factorization_machines_mnist.ipynb at master · awslabs/amazon-sagemaker-examples

事前準備

モデルを保存する場所や学習時に使用するIAMロールを定義しておきます。

bucket = '<your_s3_bucket_name_here>'
prefix = 'sagemaker/DEMO-fm-mnist'

import boto3
import re
from sagemaker import get_execution_role

session = boto3.Session(profile_name='ml_pr', region_name='ap-northeast-1')

# IAMロールのARNを取得/定義する
# role = get_execution_role() # ノートブックインスタンスから実行の場合はこれをコメントアウトして使用する
role = 'arn:aws:iam::321486094060:role/service-role/AmazonSageMaker-ExecutionRole-20190709T095048'

データ準備

MNISTの手書き数字のデータセットをダウンロードし、展開します。データセットには学習と検証、テスト用がありますが、今回は検証用を使って進めます。元々は0~9の数字のどれかを予測する分類問題ですが、今回のモデルは0かどうかの二値分類のモデルなので、ラベルを書き換えます。

import pickle, gzip, numpy, urllib.request, json, io
import sagemaker.amazon.common as smac

# データセットを読み込みます
urllib.request.urlretrieve("http://deeplearning.net/data/mnist/mnist.pkl.gz", "mnist.pkl.gz")
with gzip.open('mnist.pkl.gz', 'rb') as f:
    train_set, valid_set, test_set = pickle.load(f, encoding='latin1')

vectors = train_set[0].astype('float32')

# 0かどうかの二値分類にする
labels = (train_set[1] == 0).astype('float32')


# 学習に適したデータ形式(RecordIO)に変換する
buf = io.BytesIO()
smac.write_numpy_to_dense_tensor(buf, vectors, labels)
buf.seek(0)

学習用データをS3にアップロードします。

import boto3
import os

key = 'recordio-pb-data'
session.resource('s3').Bucket(bucket).Object(os.path.join(prefix, 'train', key)).upload_fileobj(buf)
s3_train_data = 's3://{}/{}/train/{}'.format(bucket, prefix, key)
print('uploaded training data location: {}'.format(s3_train_data))

学習

マネージドスポットトレーニングを有効化してFMモデルを学習させるトレーニングジョブを作成します。関連するパラメータは次の通りです。

(必須)train_use_spot_instances
マネージドスポットトレーニングを有効化するかどうか
デフォルト: False
(必須)train_max_wait
スポットインスタンスをプロビジョニングするまでの待機時間を含めたジョブの最大実行時間
train_max_run以上の値を設定する必要があります。
デフォルト: None
train_max_run
ジョブの最大実行時間
マネージドスポットトレーニング用のパラメータではないが、train_max_waitがtrain_max_run以上の値を設定する必要があるため、train_max_runと合わせて設定する必要があります
デフォルト: 24 * 60 * 60
checkpoint_s3_uri
ジョブ中断時のチェックポイントファイルの保存場所
デフォルト: None
checkpoint_local_path
トレーニング中にチェックポイントファイルが保存されるパス
デフォルト: /opt/ml/checkpoints/

学習処理は次のような感じになります。

from sagemaker.amazon.amazon_estimator import get_image_uri
import boto3
import sagemaker
from os import path

container = get_image_uri(session.region_name, 'factorization-machines')
output_location = path.join('s3://', bucket, prefix, 'output')
sess = sagemaker.Session(session)

fm = sagemaker.estimator.Estimator(container,
                                   role,
                                   train_instance_count=1,
                                   train_instance_type='ml.c4.xlarge',
                                   output_path=output_location,
                                   train_max_run=3600, # トレーニングジョブの最大実行時間
                                   train_use_spot_instances=True, # マネージドスポットトレーニングを有効化
                                   train_max_wait=3600, # トレーニングジョブの最大待機時間 (>=train_max_run)
                                   sagemaker_session=sess
                                  )
fm.set_hyperparameters(feature_dim=784,
                      predictor_type='binary_classifier',
                      mini_batch_size=200,
                      num_factors=10,
                      epochs=100)

fm.fit({'train': s3_train_data})

5分ほどでジョブが完了し、トレーニングログの最後にインスタンス費用の課金時間(Billable Seconds)等の情報が表示されます。

さいごに

マネージドスポットトレーニングを有効化することで、トレーニングジョブにかかるコストをかなり節約できそうです。ハイパーパラメータの調整ジョブでも同様にマネージドスポットトレーニングを有効化することができます。ハイパーパラメータの調整ジョブの場合はトレーニングジョブを複数実行するため、節約効果はさらに大きくなるかと思います。チェックポイントを有効化することでジョブの中断によるリスクをある程度は抑えられそうですし、リスクを許容できる用途で積極的に活用していきたいところです。