🕶️

AWS Sagemaker Studioで独自Containerを使う-後編-

2022/08/30に公開

AWS

はじめに

この記事はAWS Sagemaker Studioで独自Containerを使う-前編-とAWS Sagemaker Studioで独自Containerを使う-中編-の続きです。こちらを読んでいない方はそちらからどうぞ。

学習スクリプト

さて、前の記事でDockerfileのビルドを行いました。ここからそのコンテナを使用して学習を行います。
初めに入力データや出力先、containerイメージを指定します。

train.ipynd

import boto3
import sagemaker
from sagemaker import get_execution_role
from sagemaker import analytics

# setting role and session
role = get_execution_role()
sess = sagemaker.Session()
sm = boto3.Session().client('sagemaker')
region = boto3.session.Session().region_name
account = sess.boto_session.client('sts').get_caller_identity()['Account']
print(f'AccountID: {account}, Region: {region}, Role: {role}')

# input data
prefix = "dataA"
s3_input_path = f's3://your-data-bucket/{prefix}'
s3_input_data = sagemaker.inputs.TrainingInput(s3_data=s3_input_path)

# output folder
bucket_name = sess.default_bucket() 
s3_output_path = f"s3://{bucket_name}/{prefix}/build"

# instance type
instance_type = 'ml.g4dn.xlarge'

print(f'Bucket name: {bucket_name}, prefix: {prefix}, Instance type: {instance_type}')

# ECR container image
img_uri = f'{account}.dkr.ecr.{region}.amazonaws.com/classification'

input dataはS３内のデータフォルダを選択します。　IAMでアクセス許可済みのデータフォルダを指定してください。outputはモデルなどの出力先を設定します。
img_uriはECRのコンテナーイメージを選択します。学習に使用するインスタンスはml.g4dn.xlargeにしていますがここも適宜変更してください。

次にsagemaker experimentの設定を行います。experimentsとは機械学習モデルの再現性を担保するために必要な情報（モデルのバージョン管理や学習の追跡・比較・評価など）を収集して管理することができる機能で，記録をGUI上から確認することができるのでとても便利です。

train.ipynd

import time
import datetime
from smexperiments.experiment import Experiment
from smexperiments.tracker import Tracker
from smexperiments.trial import Trial
from smexperiments.trial_component import TrialComponent
from sagemaker.analytics import ExperimentAnalytics

#experiments
experiment_name = "dataA-training"
experiment = Experiment.create(experiment_name=experiment_name, 
　　　　　　　　　　　　　　　　　　　　　　　　　　description="classification for dataA.",
	　　　　　　　　　　sagemaker_boto_client=sm)
print(experiment)

次に学習時に使用するパラメータを定義します。
デープラーニング学習時は色々なパラメータを試して、最適なものを選ぶというプロセスが多いと思います。それに対応すべく、パラメータの種類と組み合わせを全て網羅し、最終的な引数を作成できるように工夫します。
下記の例ではseedの値を変更して、1と2の場合を学習してみるといった具合です。

train.ipynd

import itertools
hyperparam_options = {"batch-size": [8],
                      "epochs": [20], 
                      "lr": [1.0],
                      "gamma": [0.7],
                      "seed": [1,2],
                     }

hypnames, hypvalues = zip(*hyperparam_options.items())
trial_hyperparameter_set = [dict(zip(hypnames, h)) for h in itertools.product(*hypvalues)]
print(trial_hyperparameter_set)

最後にこれまで設定したものを含めてトレーニングさせてみましょう。

train.ipynd

run_number = 1
for trial_hyp in trial_hyperparameter_set:
    # Combine static hyperparameters and trial specific hyperparameters
    hyperparams = trial_hyp

    # Create unique job name with hyperparameter and time
    now = datetime.datetime.now()
    time_append = now.strftime('%Y%m%d-%H%M')

    hyp_val = list(trial_hyp.values())
    hyp_append = "-".join([str(elm).replace(".", "-") for elm in hyp_val])
    print(hyp_append)
    training_job_name = f"job-{hyp_append}-{time_append}"
    trial_name = f"trial-{hyp_append}-{time_append}"
    trial_desc = f"training-{run_number}"
    training_tags =[{"Key": "trial-desc", "Value": trial_desc}]
    
    # Create a new Trial and associate Tracker to it
    experiments_trial = Trial.create(
        trial_name=trial_name,
        experiment_name=experiment.experiment_name,
        sagemaker_boto_client=sm,
        tags=training_tags,
    )

    # Create an experiment config that associates training job to the Trial
    experiment_config = {
        "ExperimentName": experiment.experiment_name,
        "TrialName": experiments_trial.trial_name,
        "TrialComponentDisplayName": training_job_name,
    }

    metric_definitions=[
        {'Name': 'epoch', 'Regex': 'Epoch: (.*?),'},
        {'Name': 'training-loss', 'Regex': 'Training Loss: (.*?),'},
        {'Name': 'test-loss', 'Regex': 'Test Loss: (.*?),'},
        {'Name': 'test-acc', 'Regex': 'Test Acc: (.*?)%;'},
    ]
    
   # Estimator
    estimator = sagemaker.estimator.Estimator(
        image_uri=img_uri,
        role=role,
        instance_count=instance_count,
        instance_type=instance_type,
        hyperparameters=hyperparams,
        output_path=s3_output_path,
        sagemaker_session = sess,
        base_job_name="job",
        metric_definitions=metric_definitions,
        enable_sagemaker_metrics=True,
        tags=training_tags,
    )


    # Fit Launch a training job
    estimator.fit(
        {'train': s3_input_data},
        job_name=training_job_name,
        experiment_config=experiment_config,
        wait=True,
    )

    # give it a while before dispatching the next training job
    time.sleep(2)
    run_number = run_number + 1

metric_definitionsではトレーニングスクリプトで出力されるlogをキャッチして保存します。

estimaterでは入力データをやトレーニング名、experimentsのconfigを渡します。
wait=Falseにすればバックグラウンドで実行され、結果は出力されません。
また、複数のjobを実行する場合も前のジョブの完了を待たずに実行できます。

    estimator.fit(
        {'train': s3_input_data},
        job_name=training_job_name,
        experiment_config=experiment_config,
        wait=True,
    )

結果確認

無事トレーニングが終わったら結果を確認してみましょう。studioの左下のアイコンを押して、「Experiments and trials」を選択します。

下に表示されるリストから、先ほど行ったトレーニングを選択してみましょう。
Chartに下記のような学習履歴が表示されればOKです。

experimentsを使用することで、このようにGUIコンソールから結果をいつでも確認することができます。あとは、公式にあるように、モデルをデプロイするなりモニタリングするなり好きに運用してみましょう。

まとめ

今回はsagenmaker studioで独自コンテナーをビルドして、トレーニングを行ってみました。
どうでしょうか、一部煩わしい設定等はあるものの、一度構築してしまえばボタンひとつで学習が行える手軽さが魅力です。AIモデルがあらゆるサービスに使われるようになった今日、特にAWSでサービスを運用している方には非常に重要な技術ではないでしょうか。

それではまた！

はじめに

学習スクリプト

結果確認

まとめ

Discussion