【アドベントカレンダーDay 2 】SageMaker AIで画像分類モデルをFine-tuningする手順

先日、Jr. ChampionsのイベントでFine-tuningのハンズオンを行いました。本記事では、そのハンズオンの内容をベースに、Amazon SageMaker AIを使って画像分類モデル（ResNet18）をFine-tuningする手順を解説します。

「自分と他人を識別するモデル」を作成することを目標に、事前学習済みモデルの重みを活用して少ないデータでも高精度なモデルを構築する方法を紹介します。

Fine-tuningとは

Fine-tuningは、別のタスクで学習した知識を応用して新しいタスクを学習する手法です。具体的には、事前に学習済みの重みから学習を開始します。

メリットは以下の通りです。

• すでに学習された知識を最初から使えるので時間短縮になる
• 応用できる範囲が広い（例：1000クラス分類モデル → 2クラス分類モデル）

今回は、ImageNetで1000クラス分類を学習済みのResNet18を、「自分（me）」と「他人（notme）」の2クラス分類にFine-tuningします。

本記事で使用する環境

項目	内容
SageMaker Studio	Code Editor
トレーニングインスタンス	ml.p3.2xlarge（NVIDIA Tesla V100）
フレームワーク	PyTorch 2.3
Pythonバージョン	3.11
ベースモデル	ResNet18（ImageNet事前学習済み）

事前準備：SageMaker AIドメインの作成

SageMaker AIを使用するには、ドメインの作成が必要です。AWSコンソールから「Amazon SageMaker AI」を検索し、「シングルユーザー向けの設定」からドメインを作成します。

ドメイン作成後、ユーザープロファイルからStudioを起動し、Code Editorを選択します。Code EditorはVSCodeライクなIDEで、Pythonがデフォルトで使用可能です。

Code Editor Spaceの作成

1. Code Editorから「Create Code Editor space」を選択
2. Name: handson-space として作成
3. 作成完了後、「Run space」でSpaceを起動
4. 「Open Code Editor」でエディタを開く

ホームディレクトリは /home/sagemaker-user/ です。

データセットの準備

フォルダ構成

以下のフォルダ構成でデータを格納します。

data/
├── train/
│   ├── me/          # 自分の顔画像（トレーニング用）
│   └── notme/       # 他人の顔画像（トレーニング用）
└── test/
    ├── me/          # 自分の顔画像（テスト用：各4枚）
    └── notme/       # 他人の顔画像（テスト用：各4枚）

データセット作成のヒント

良いモデルを作るためのポイントです。

• 明るさや暗さが異なる画像を含める
• メガネの有無など、バリエーションを持たせる
• 他人の画像には異性も含める（テストデータに含まれる場合）

テスト用には各クラス4枚ずつ選び、残りをトレーニング用にします。

トレーニングの実装

main.ipynb の作成

main.ipynb というNotebookファイルを作成し、以下のコードを順番に実行していきます。

SageMaker SDKのインポート

# sagemakerをPythonから操作するライブラリをインストールします。
import sagemaker

# sagemakerのセッションを取得します
session = sagemaker.Session()

# sagemakerの実行ロールを取得します
execution_role = sagemaker.get_execution_role()
print(execution_role)

トレーニングデータをS3にアップロード

# トレーニングに使用するデータを管理します。

# ワークスペースのトレーニングデータが置いてあるパス
data_dir = "data/train"
# トレーニングデータをおくS3のprefix
s3_prefix = "image-training-data"

# S3にトレーニングデータをアップロードします
s3_train_path = session.upload_data(path=data_dir, key_prefix=s3_prefix)
print(f"データをS3にアップロードしました: {s3_train_path}")

train.py の作成

トレーニングジョブで実行するスクリプトを作成します。

# 入力の引数をパースするためのライブラリ
import argparse
# OSに関する操作の関数が入ったライブラリ
import os
# 機械学習用のライブラリ
import torch
# torchの中でも画像処理タスクに特化したライブラリ
from torchvision import datasets, transforms, models

# 引数をうけとります
parser = argparse.ArgumentParser()
# エポック数（全データセットを何周学習するか）
# - 小さすぎる: 学習不足 (underfitting)
# - 大きすぎる: 過学習 (overfitting)
parser.add_argument("--epochs", type=int, default=5)

# バッチサイズ（一度に何枚の画像をまとめて処理するか）
# - 小さすぎる: 学習が遅い、学習が不安定
# - 大きすぎる: メモリを大量に使う、局所最適解にはまりやすい
parser.add_argument("--batch_size", type=int, default=8)

# 学習率（モデルの重みをどれくらいの大きさで更新するか）
# - 小さすぎる: 学習が遅い
# - 大きすぎる: 最適解を飛び越えてしまう
parser.add_argument("--learning_rate", type=float, default=1e-4)

# トレーニングデータへのパス
parser.add_argument("--train", type=str, default="/opt/ml/input/data/train")

# 引数を取得します。
args = parser.parse_args()
# これで、たとえば arg.train で、トレーニングデータへのパスにアクセスできます。


# モデルを定義します。
# weights で重みを指定します。Fine-tuningなので、トレーニング済みの重みを指定します。
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# 再学習をする層を最後の総結合層のみにします。
# 層を「凍結させる」と言ったりします。
for name, param in model.named_parameters():
    if "fc" in name:
        param.requires_grad = True
    else:
        param.requires_grad = False

# モデルの出力を2つにします。最後の層（総結合層）を出力が2の層に置き換えます。
model.fc = torch.nn.Linear(model.fc.in_features, 2)

# 入力画像を前処理します。
# モデルの入力にサイズが合うように 224 x 224 にResizeします。
# モデルの入力に形式が合うように、テンソルに変換します。
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])


# トレーニング用のデータを取得します。
# transformを指定することで、取得したデータに自動でpreprocessを施します。
train_dataset = datasets.ImageFolder(args.train, transform=preprocess)
# トレーニングがしやすいようにモデルをロードします。
# バッチサイズに合わせて、毎エポックで画像をシャッフルしてくれるようにします。
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)


# GPUの設定をします。
device = torch.device("cuda")
# モデルをGPUに転送します。
model = model.to(device)


# トレーニングアルゴリズムを定義します。
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)


# トレーニングをします。
# 全体をエポックの数だけ学習します。
for epoch in range(args.epochs):
    # モデルを学習モードに設定します。
    model.train()
    
    # このエポックでの損失の合計を記録します。
    # 直接トレーニングには必要ありませんが、トレーニングが適切に進んでいるかを判断するために必要です。
    running_loss = 0.0
    
    # バッチごとにデータを取り出して学習します。
    for imgs, labels in train_loader:
        # 画像とラベルをGPUに転送します。
        imgs, labels = imgs.to(device), labels.to(device)
        
        # 前回の勾配をリセットして、今回のトレーニングに備えます。
        optimizer.zero_grad()
        # 現状のモデルで予測を行います。
        outputs = model(imgs)
        # 予測と正解の差（損失）を計算します。
        loss = criterion(outputs, labels)
        
        # 損失から勾配を計算します。
        loss.backward()
        # 勾配を使って、モデルのパラメータを更新します。
        optimizer.step()
        
        # このバッチでの損失を記録します。
        running_loss += loss.item()
    
    # エポックごとの平均損失を表示します。
    print(f"エポック [{epoch + 1}/{args.epochs}]: ロス {running_loss / len(train_loader):.4f}")


# 学習し終わったモデルを保存します。
# /opt/ml/model ディレクトリを作成します。
os.makedirs("/opt/ml/model", exist_ok=True)
# /opt/ml/model に置いたファイルはトレーニング終了後に自動でS3に保存されます。
torch.save(model.state_dict(), "/opt/ml/model/model.pth")
print("モデルは保存されました。")

トレーニングジョブの実行

main.ipynb に戻り、トレーニングジョブを開始します。

# トレーニング用のライブラリのインポート
from sagemaker.pytorch import PyTorch

# トレーニングジョブを定義します。
job = PyTorch(
    entry_point="train.py",
    role=execution_role,
    framework_version="2.3",
    py_version="py311",
    instance_type="ml.p3.2xlarge",
    instance_count=1,
    hyperparameters={
        "epochs": 15,
        "batch_size": 6,
        "learning_rate": 5e-4
    },
)

# トレーニングジョブを開始します。
job.fit({"train": s3_train_path})

トレーニングジョブの進行状況は、SageMaker AIコンソールの「トレーニングジョブ」から確認できます。「ログを表示」をクリックすると、CloudWatch Logsでエポックごとの損失を確認できます。

推論の実行

モデルのダウンロード

# モデルを出力したS3へのファイルパスを取得します。
model_artifact = job.model_data
print(f"モデルが保存されている場所: {model_artifact}")

# モデルをダウンロードします。
!aws s3 cp {model_artifact} ./model.tar.gz

# 圧縮されたダウンロードしたモデルを解凍します。
!tar -xzvf model.tar.gz

推論関数の定義

import torch
from torchvision import models, transforms
from PIL import Image
import glob
import os
import time

# モデルを定義します。
# トレーニング時と同じ構造であることが重要です。
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.fc = torch.nn.Linear(model.fc.in_features, 2)

# モデルの重みをロードします。
model.load_state_dict(torch.load("model.pth", map_location=torch.device("cpu")))
# 推論モードに切り替えます。
model.eval()

# 推論用の画像の前処理をします。
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

def predict_and_display(img_dir):
    """
    指定されたディレクトリ内の画像を一枚ずつ推論し、画像と結果を表示します。
    Args:
        img_dir (str): 画像が格納されているディレクトリのパス
    """
    # 画像の名前を取得します。
    img_files = glob.glob(os.path.join(img_dir, "*.*"))
    # 画像を一枚ずつ順に見ていきます。
    for img_file in img_files:
        # 画像を取得します。
        img = Image.open(img_file)
        # 表示用の画像です。
        img_display = img.copy()
        # サイズを 224 x 224 にします。画像のサイズが大きすぎると Notebook で表示できないからです。
        img_display.thumbnail((224, 224), Image.Resampling.LANCZOS)
        # 画像を表示します。
        img_display.show()
        
        # 画像をモデルに適した形式に変換します。（224 x 224 x 3 のテンソル）
        input_tensor = preprocess(img)
        # 画像を一枚だけからなるバッチに変換します。
        input_batch = input_tensor.unsqueeze(0)
        # 推論をします。このとき勾配の計算は必要ないので行いません。
        with torch.no_grad():
            output = model(input_batch)
        # me, notme の確率を計算します。つまり、0以上1以下の値になるように変換します。
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        # 予測されたクラスを計算します。0ならme、1ならnotmeです。
        predicted_class = probabilities.argmax().item()
        # 予測されたクラスである確率を計算します。
        confidence = probabilities.max().item()
        # 予測結果を表示します。
        print(f"ファイル名: {img_file}")
        print(f"推論結果: {'Me' if predicted_class == 0 else 'Not Me'}")
        print(f"自信度: {confidence:.2f}")
        print("-" * 30)
        
        time.sleep(3)

推論の実行

# meフォルダを推論します
predict_and_display("data/test/me")

# notmeフォルダを推論します
predict_and_display("data/test/notme")

ハイパーパラメータの選び方

トレーニング結果を見て、以下の点を確認します。

現象	原因	対策
損失が下がりきらない	学習不足（underfitting）	エポック数を増やす、学習率を上げる
損失が途中から上がり始める	過学習（overfitting）	エポック数を減らす、データを増やす
損失が振動して安定しない	学習率が高すぎる	学習率を下げる
損失が全く下がらない	学習率が低すぎる	学習率を上げる

今後の応用

本記事ではCode Editor上で推論を行いましたが、実運用ではSageMakerエンドポイントとしてデプロイすることが考えられます。また、より精度を上げるためにデータ拡張（Data Augmentation）を適用することも有効です。これは今後機会があれば書いていきたいです。

まとめ

本記事では、Amazon SageMaker AIを使って画像分類モデル（ResNet18）をFine-tuningする一連の手順を解説しました。

特に以下の点を紹介しました。

• Fine-tuningの概念と利点
• SageMaker Studio Code Editorの使い方
• PyTorchを使ったトレーニングスクリプトの実装
• SageMakerトレーニングジョブの実行方法
• 学習済みモデルを使った推論

Fine-tuningを使うことで、少ないデータでも高精度なモデルを短時間で構築できます。画像認識だけでなく、LLMのFine-tuningも基礎的な構造は同じです。「既存のモデルでは物足りない」という場面で、Fine-tuningを検討してみてください。

時間がなかったので割愛したんですが、SageMaker Endpointにデプロイするなどもできると良さそうでした。また、オリジナルの画像モデルは流行りのAIエージェントなどでも個別特化タスクをするのにとても重宝します。AgentCore Gatewayを使ってオリジナルモデルをエージェントから使用するなどの最近流行りの応用もあります。いつか記事として公開できるといいなと思います。