SageMakerとTensorFlow+Kerasで始めるDeep Learning

はじめに

ここでは、

• TensorfFlow+Kerasを用いて画像認識モデルを自作するための基本的な知識を習得する
• クラウドサービス上でモデルを構築・利用するイメージを掴む

の2点を目的として、Amazon SageMaker上で簡単な深層学習モデル（画像認識モデル）を作成した際の手順をまとめました。

また、今回はre:Invent 2024で発表されたAmazon SageMaker Unified Studioを利用してみました。（基本的なモデル構築しか行っていないので、従来のSageMakerからの進化はあまり活かせていませんが、試しに利用してみた次第です。）

SageMakerとTensorflow+Kerasの環境構築

SageMaker Unified Studio

今回はAIモデルの開発にAmazon SageMakerを、中でもSageMaker Unified Studio(プレビュー)を利用しました。

ここで、SageMaker Unifide Studioは

• すべてのデータと AI を統合したエクスペリエンス
• どんなジョブにも、最高クラスのツールを使う
• AI モデルを大規模にトレーニング、カスタマイズ、デプロイ
• カスタムの生成 AI アプリケーションを迅速に構築
• Amazon Q Developer でデータジャーニーを加速
  (Amazon SageMaker Unified Studio)

を特徴としており、従来のSageMakerではAIモデルの開発・管理・デプロイを中心に扱っていたのに対し、新たにデータの取り込みから分析、最新の生成AI（基盤モデル）の活用まで、より広範なAIプロジェクトのライフサイクル全体を1つの環境で管理できるように進化したものである。と言えるようです。

TensorFlow＋Keras

今回はPython上で機械学習ライブラリとしてTensorFlowおよびKerasを利用します。
これらはGoogleが開発した機械学習フレームワークで、ニューラルネットワークの構築においてPytorchなどと並んで広く用いられています。

それぞれのイメージとしては以下になります。

• TensorFlow: ニューラルネットワークの学習に必須の分散処理やGPU計算、微分演算や最適化などの計算を効率的に処理するためのフレームワーク
• Keras: Tensorflow上でのモデル構成や最適化の方法などをPythonで簡単に記述するためのフレームワーク

なお、TensorFlow 2.x.x以降ではKerasが統合されており、TensorFlowからKerasを呼び出し可能になっています。

環境構築

以下に、SageMaker Unified Studioを利用するための環境構築手順を記載します。
今回は、AWSの提供する専用のCloudFormationを用いて環境の構築をしました。

①：SageMakerドメインの作成

サービスとしてSageMakerを選択し、ドメインを新規に作成します。ここで、ドメインとは個別の開発環境を表し、マネジメントコンソールとは別にSageMaker専用の管理画面が作成されます。

作成する際に、VPCやS3等と紐付ける必要がありますが、VPCを作成をクリックするとそれらの関連リソースのネットワークを確立するCloudFormationを利用することができます。（VPC内に3つのプライベートサブネット、1つのパブリックサブネット、S3接続用のVPCエンドポイントが作成されます。）

作成されたVPCおよびサブネットを指定してドメインの作成を完了します。（今回はインターネット経由でのモデルアクセスを行わないので、プライベートサブネットのみ選択しました。）

なお、今回SageMakerの利用はIAMユーザーとしてログインして行なっており、該当IAMユーザーには権限としてAdministratorAccessを与えています。

②：Tensorflow+Kerasの開発環境作成

Pythonを用いて機械学習モデルを構築するにあたり使用する環境は、SageMakerを立ち上げた時点で概ね整っているので特別な作業は不要となります。
ただし、SageMakerにおけるコンピューティングインスタンスの構成は理解しておくと今後の作業がしやすいため、こちらで整理します。

まず、マネジメントコンソール上のドメイン一覧からSageMakerの開発環境にアクセスします。
SageMaker上ではプロジェクト単位で開発環境やデータ、作成モデル等を管理するため、新規にプロジェクトを作成します。今回、名前以外はデフォルトのままで作成しており、犬と猫の画像分類を行ったためプロジェクト名は"Dog-Cat"としました。

今回はPython3を用いてモデル構築〜デプロイを行うので、JupyterLabを利用します。

JupyterLabを開くと以下の画面のようにワークスペースが表示され、ここで作業を行うことになります。

今回は、大きく2つのPythonスクリプトをしようすることになるのですが、それぞれのスクリプトの実行リソースが異なる（＝実行環境が異なる）ので注意が必要です。
具体的には、以下の図のように3種類のインスタンスを用いることになります。

• ノートブックインスタンス = コードを開発用 + SageMakerのインタフェイス
• 学習用インスタンス = 作成したコードを用いてモデルの学習を実行
• 推論用インスタンス = 作成したモデルのデプロイ先

Amazon SageMaker Studio で機械学習の最初の一歩を踏み出そう より

ただし、いずれの環境もTensorFlowを新たにinstallする必要はなく基本的なライブラリも利用可能なので、今回は新規のライブラリのインストール等は行わずに進めます。

モデル構築

JupyterLab上でモデルを構築していきます。前述の通り、今回は画像認識の初歩として犬猫の画像分類モデルを構築します。モデルの大枠としては畳み込みニューラルネットワーク（CNN）を利用し、事前学習なしのフルスクラッチ開発とします。

データ準備

機械学習には多くのデータが必要となりますが、今回はKaggleのCat and Dogというデータセットを利用します。
以下リンクから画像のダウンロードを行いました。

• https://www.kaggle.com/datasets/tongpython/cat-and-dog

データセットの内訳としては以下のテーブルの通りです。本来は学習データの内何割かを検証データ（ハイパーパラメータ調整用のデータ）とし、チューニングしたモデルをテストデータで評価するのですが、今回はモデルの評価まで行わないためテストとして分けられているデータを検証データとして扱います。

	Training	Test
Cat	4000	1011
Dog	4005	1012

以上の画像データを学習に用いるにあたり、S3上にデータを配置する方法と、ノートブックインスタンスに直接データを配置する方法がありますが、今回はS3上のデータを利用する方法を選択します。従って、ダウンロードした画像をS3上へアップロードしていきます。ここで格納先のS3バケットはSageMaker作成時にデフォルトで紐づけられたものとします。この対象バケットが不明の場合はnotebook上で以下を実行して確認可能です。

S3Bucket.py

import sagemaker

session = sagemaker.Session()
default_bucket = session.default_bucket()
print(f"Default bucket: {default_bucket}")

該当バケット内にプロジェクトのフォルダが存在するため、その配下にdatasetというフォルダを作成しデータを格納します。

以上でデータの準備は完了です。

モデルの作成

いよいよモデルの構築および学習を行っていきます。
モデルの構築および学習におけるコーディングは以下を参考に行っています。

• Amazon SageMakerの基礎
• Amazon SageMakerで簡単にKerasを使う方法

モデルの構築

JupyterLab上にモデルの構築を行うPythonファイルcnn_sagemaker_tensorflow.pyを/src下に作成します。（このファイルの実行は学習用インスタンスで行われます。）

コード概要

• ハイパーパラメータはノートブックインスタンス側から受け取り
• モデルやデータの格納先はSageMakerの環境変数として利用可能
• 画像データの前処理はKerasのImageDataGeneratorで実施（データ拡張は未適用）
• モデルのアーキテクチャはtf.keras.Sequential内で定義
• mlflowで実験データを保管（詳細は後述）

cnn_sagemaker_tensorflow.py

import numpy as np
import argparse
import os
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras import layers
import mlflow

def train(args):
    # バックエンドの設定
    K.set_image_data_format('channels_last')

    batch_size = args.batch_size
    epochs = args.epochs
    num_classes = args.num_classes
    train_dir = args.train
    test_dir = args.test
    learning_rate = args.learning_rate
    dropout_rate = args.dropout_rate

    # Model / data parameters
    input_shape = (128, 128, 3)
    image_size = (128, 128)

    # データの準備
    train_datagen = ImageDataGenerator(
        rescale=1./255,
    )

    test_datagen = ImageDataGenerator(
        rescale=1./255,
    )
    
    train_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=image_size,
        batch_size=batch_size,
        class_mode='categorical'
    )
    
    test_generator = test_datagen.flow_from_directory(
        test_dir,
        target_size=image_size,
        batch_size=batch_size,
        class_mode='categorical'
    )

    # モデルのアーキテクチャを定義
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(128, kernel_size=(3, 3), activation="relu", padding='same'),
        layers.MaxPooling2D(pool_size=(2, 2)),
        # layers.Conv2D(256, kernel_size=(3, 3), activation="relu", padding='same'),
        # layers.MaxPooling2D(pool_size=(2, 2)),
        layers.GlobalAveragePooling2D(),
        layers.Dropout(dropout_rate),
        layers.Dense(32, activation="relu"),
        layers.Dropout(dropout_rate),
        layers.Dense(num_classes, activation="softmax")
    ])
    
    model.summary()

    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
    model.compile(
        loss="categorical_crossentropy", 
        optimizer=optimizer,
        metrics=["accuracy"]
    )
    
    
    # 環境変数から設定を取得
    tracking_uri = os.getenv('MLFLOW_TRACKING_URI')
    experiment_name = os.getenv('MLFLOW_EXPERIMENT_NAME')
    experiment_id = os.getenv('MLFLOW_EXPERIMENT_ID')
    run_name = os.getenv('MLFLOW_RUN_NAME', 'default-run')
    
    print(f"MLflow Tracking URI: {tracking_uri}")
    print(f"MLflow Experiment Name: {experiment_name}")
    print(f"MLflow Experiment ID: {experiment_id}")
    
    # MLflow設定
    if tracking_uri:
        mlflow.set_tracking_uri(tracking_uri)
    
    if experiment_name:
        mlflow.set_experiment(experiment_name)
    
    with mlflow.start_run(run_name=run_name) as run:
        print(f"Started MLflow run: {run_name} (ID: {run.info.run_id})")

        # モデルの訓練
        history = model.fit(
            train_generator,
            validation_data=test_generator,
            epochs=epochs,
            steps_per_epoch=train_generator.samples // train_generator.batch_size,
            validation_steps=test_generator.samples // test_generator.batch_size
        )
        
        for epoch in range(len(history.history['loss'])):
            mlflow.log_metric("training_loss", history.history['loss'][epoch], step=epoch)
            mlflow.log_metric("val_loss", history.history['val_loss'][epoch], step=epoch)
            mlflow.log_metric("training_accuracy", history.history['accuracy'][epoch], step=epoch)
            mlflow.log_metric("val_accuracy", history.history['val_accuracy'][epoch], step=epoch)

    # モデルの保存
    model_path = os.path.join(args.model_dir, 'model/1')
    model.save(model_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # hyperparameterはSageMakerから変数として受け取り
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--num_classes', type=int, default=2)
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--learning_rate', type=float, default=1e-3)
    parser.add_argument('--dropout_rate', type=float, default=0.5)

    # データとモデルの格納先は環境変数を使用
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
    parser.add_argument('--test', type=str, default=os.environ['SM_CHANNEL_TEST'])
    
    args, _ = parser.parse_known_args()
    train(args)

モデルの学習

モデルの学習を行うにあたり、SageMakerでは実験管理ツールとしてフルマネージドのMLflowを利用することができるため、これを用いてモデル作成の実験管理（学習時のハイパーパラメータや学習過程、性能の記録）を行っていきます。

注意
ただし、MLflowの利用にはMLflow Tracking Serverを設置する必要があり、1時間あたりUSD 0.789の課金が発生します(参考)。実験管理の使用感を理解したら、すぐにインスタンスを削除することを推奨します。

Unified Studioからの利用は、プロジェクトのコンピューティングからMLflow Tracking Serverタブ内で新規にサーバーを追加することで可能です。

では、cnn_sagemaker_tensorflow.pyを用いて学習を実行するためのスクリプトを/src下のJupyter notebook上で実行していきます。

①まずはS3上に保存した学習データを指定

• 先ほどデータを保管したS3内の学習データ、テストデータの保存先フォルダを指定

SageMaker-1.py

import os
import sagemaker

os.makedirs("./data", exist_ok = True)

sagemaker_session = sagemaker.Session()
bucket_name = sagemaker_session.default_bucket()

#S3上のフォルダ構成を記載
training_data = f's3://{bucket_name}/dzd_aplszek1ws0st5/auws0i5t6r0kux/dev/dataset/training_set/training_set'
test_data = f's3://{bucket_name}/dzd_aplszek1ws0st5/auws0i5t6r0kux/dev/dataset/test_set/test_set'


training_input = sagemaker.inputs.TrainingInput(
    s3_data=training_data,
    content_type='application/x-image'
)

test_input = sagemaker.inputs.TrainingInput(
    s3_data=test_data,
    content_type='application/x-image'
)


print('Training data is in: {}'.format(training_data))
print('Validation data is in: {}'.format(test_data))

②SageMakerとMLflowの基本設定

• MLflow tracking serverのARNをコピーしてmlflow_tracking_arnとして指定
• experiment_nameおよびrun_nameは実験データを区別できるように名付ける

SageMaker-2.py

import mlflow
import boto3
from sagemaker.tensorflow import TensorFlow
from sagemaker import get_execution_role
from mlflow.tracking import MlflowClient

# SageMaker基本設定
sagemaker_session = sagemaker.Session()
role = get_execution_role()
region = sagemaker_session.boto_region_name

# MLflow Tracking Serverの設定
mlflow_tracking_arn = '<MLflow tracking serverのARN>'
mlflow.set_tracking_arn(mlflow_tracking_arn)

# 実験の設定
experiment_name = "dog_cat_1"
mlflow.set_experiment(experiment_name)
run_name = "run-1" 

③学習の実行

• ハイパーパラメータを指定
• estimatorの定義において
  • entry_point = ./cnn_sagemaker_tensorflow.pyでモデル構築を定義したファイルを指定
  • input_mode='File'でS3上の学習データの取り込み方式を指定（参考）
    • ファイルモードでははじめに全ファイルをダウンロードしてから学習
  • 学習用インスタンスで必要なライブラリがインストールされていない場合はdependencies=['requirements.txt']で指定するrequirements.txtに記載し、txtファイルを同階層に配置する
    ./requirements.txt

    mlflow==2.13.2
    sagemaker-mlflow==0.1.0
    

• estimator.fitで学習用インスタンスを立ち上げて学習実行
  • 'instance_type': "ml.m5.xlarge"で立ち上げるインスタンスタイプを指定
    • 学習終了と同時にインスタンスは終了する（今回の学習はm5.xlargeで60分程度）

SageMaker-3.py

#hyperparameterの設定
epochs=30
learning_rate=1e-4
batch_size=32
dropout_rate=0.5


# TensorFlow Estimatorの設定
estimator = TensorFlow(
    entry_point="./cnn_sagemaker_tensorflow.py",
    role=role,
    train_instance_count=1,
    train_instance_type="ml.m5.xlarge",
    framework_version="2.12.0",
    py_version='py310',
    script_mode=True,
    input_mode='File',
    hyperparameters={
        'batch_size': batch_size,
        'num_classes': 2,
        'epochs': epochs,
        'learning_rate': learning_rate,
        'dropout_rate': dropout_rate
    },
    environment={
        'MLFLOW_TRACKING_URI': mlflow_tracking_arn,
        'MLFLOW_EXPERIMENT_NAME': experiment_name,
        'MLFLOW_EXPERIMENT_ID': mlflow.get_experiment_by_name(experiment_name).experiment_id,
        'MLFLOW_RUN_NAME': run_name,
        'AWS_DEFAULT_REGION': region
    },
    source_dir='.',
    dependencies=['requirements.txt']
)

# MLflowでの実験実行
with mlflow.start_run(run_name=run_name) as run:
    # アーキテクチャとハイパーパラメータのログ
    mlflow.log_param("model_architecture", "32-64-128")
    mlflow.log_params({
        'batch_size': batch_size,
        'num_classes': 2,
        'epochs': epochs,
        'learning_rate': learning_rate,
        'dropout_rate': dropout_rate,
        'instance_type': "ml.m5.xlarge",
        'framework_version': "2.12.0"
    })   
    # SageMakerでのトレーニング実行
    estimator.fit({
        'training': training_input,
        'test': test_input
    })
    
    
    
print("===Trained model was saved===")
print(f"Run ID: {run.info.run_id}")

モデルの改良

学習済みモデルについて、学習曲線をもとにアーキテクチャやハイパーパラメータを調整していきます。まずMLflowのダッシュボードへアクセスします。

実行した学習のrun_nameに対して実験記録を確認することができます。

1. 各種ハイパーパラメータや学習時間などの学習情報
   

2. 学習中の正解率と損失の変化（学習曲線）
   

学習曲線を見ると、検証データにおける損失（破線）があまり減少しておらずかつ25エポック目以降でかなり上昇してしまっていることがわかります。つまり、作成したモデルは過学習状態であると言えそうです。

そこで、モデルのチューニングを実施したいと思います。例えば、損失関数の減少の様子はそれほど極端でないのでひとまずは学習率は維持し、cnn_sagemaker_tensorflow.pyのtf.keras.Sequential部分を以下のように変更してモデルを単純化する。などが考えられます。

model_change.py

model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),
        layers.MaxPooling2D(pool_size=(2, 2)),
        #畳込み層を一層減らす
        #layers.Conv2D(128, kernel_size=(3, 3), activation="relu", padding='same'),
        #layers.MaxPooling2D(pool_size=(2, 2)),
        # layers.Conv2D(256, kernel_size=(3, 3), activation="relu", padding='same'),
        # layers.MaxPooling2D(pool_size=(2, 2)),
        layers.GlobalAveragePooling2D(),
        layers.Dropout(dropout_rate),
        layers.Dense(32, activation="relu"),
        layers.Dropout(dropout_rate),
        layers.Dense(num_classes, activation="softmax")
    ])

他にも過学習を防ぐためのアプローチとしては、データ拡張や最適化におけるweight decayの導入などがあり、いずれもKerasで簡単に実装可能なので適宜用いていくと良いでしょう。（今回は時間の都合上割愛しました。）

以上のようにコード中のパラメータやモデル構成を変更し、モデルを最適化していきます。

モデルのデプロイ

学習済みのモデルは推論用インスタンスにデプロイすることで利用可能になります。

注意
推論用インスタンスは停止処理を行うまで稼働し続けます。稼働中は料金が発生する（東京リージョンで1時間あたりUSD 0.3程度）ため、利用終了後はインスタンスの停止を行なってください。

SageMaker-4.py

#直前に学習していたモデル（estimatorで指定されているモデル）をデプロイ
predictor = estimator.deploy(instance_type='ml.m5.xlarge', initial_instance_count=1)

S3上に保存されているモデルを利用する場合はこちらのリアルタイム推論用のモデルをデプロイするを参照してください。

それでは、デプロイしたモデルを用いて推論を実施してみます。
今回はノートブック上から推論インスタンスにアクセスするため、推論対象とする入力画像を/src下に配置します。配置する画像は自身の画像フォルダ内の適当な犬や猫の画像で問題ありません。

私は以下の2枚の画像をdog.pngとcat.pngとして配置しました。

下記コードを実行すると、

SageMaker-5.py

import numpy as np
from PIL import Image

# 画像の読み込みと前処理
img = Image.open ('dog.png')
img = img.convert ('RGB')
img = img.resize((128, 128))
img_array = np.array(img)
img_array = img_array/255.0

# バッチ次元の追加
img_batch = np.expand_dims(img_array, axis=0)

response = predictor.predict(img_batch)
print(response)

以下レスポンスが得られました。
{'predictions':［［0.00814406853，0.991856113]]}

今、リストとクラスの関係はアルファベット順でresponse[0]はCatに, response[1]はDogに対応しているので、入力した犬の画像は正しく分類できていることが分かります。
cat.pngについても同様に試すとレスポンスは{'predictions': [[0.998341918, 0.00165801216］]｝となり、こちらも正しい分類ができていることが分かります。

モデルを利用できることがわかったので、推論用インスタンスを停止します。

SageMaker-6.py

predictor.delete_endpoint()

後片付け

以上でモデルの構築〜デプロイまでを試したので、作成したリソースを消去していきます。

MLflow tracking serversの停止

コンピューティングのMLflow tracking serversから停止を選択します。

プロジェクトの削除 ※インスタンス停止失敗により緊急手段として実施

※こちらの作業は他の削除処理が正常に完了した場合は必ずしも必要でない作業となります。私が試した際には、後述のSageMakerドメインの削除をしても継続してドメインが利用可能（SageMakerにログイン可能）かつMLflow tracking serversもステータスがStopFailedとなってしまっていたため、緊急手段としてプロジェクト単位の削除を実施しました。
(この作業ではSageMakerの実行ロール（デフォルト：AmazonSageMakerProvisioning~~）が利用されますが、デフォルトでは権限不足となるため、一時的にAdministratorAccessを付与して作業を実施しました。削除完了後はポリシーのアタッチを解除してください。)

Project overviewのアクションから削除を選択し、プロジェクトの削除を実施します。

SageMakerドメインの削除

SageMakerドメインを削除するためには、前段としてドメイン内の「スペース」および「ユーザープロファイル」を削除する必要があります。

該当ドメインの詳細のタブ「スペースの管理」からスペースを削除します。

続いてタブ「ユーザープロファイル」からユーザープロファイルを削除します。

最後にドメインの削除を実施します。

S3バケットの削除

S3バケットも新規に作成されているため、こちらも削除しておきます。

その他リソースの削除

はじめにCloudFormationを利用していくつかのリソースを作成したため、それらのリソースを削除します。

マネジメントコンソールからCloudFormationを選択し、今回作成されたスタックそれぞれを削除していきます。

以上で今回作成したリソースのクリーンアップは終了です。

おわりに

ここまで、SageMakerで画像認識モデルの構築を実施した結果として私自身は、

• TensorfFlow+Kerasを用いたコードベースのモデル構築の基本的なやり方を理解する

  • Kerasを利用することでかなり直感的に構築可能である
  • クラスやモジュールも多岐に渡るため応用的な実装も可能である

• クラウドサービス上でモデルを構築・利用するイメージを掴む

  • 環境構築が不要かつリソースも必要な量を必要なときに利用可能である
  • 特にグループで開発するときは環境や資材を簡単に一元管理できる

というように理解が進みました。

ただ、特にSageMakerについては今回利用したサービスはごく一部に過ぎず、Data Wranglerをはじめとするデータ準備のサービスやハイパーパラメータの最適化手法、GUIベースでのモデル構築手段など様々なサービスが提供されているため、今後それらにも触れながらよりCloud上でのAI開発を知れたらと思います。