SageMaker AIでMNISTの画像分類モデルを作ってみた
はじめに
AWS Certified AI Practitionerを勉強する中で、SageMakerについて問われる問題が多く、せっかくなら触ってみようということで簡単な内容ではありますが、
Amazon SageMaker AIを使って、手書き数字のデータセットMNISTの画像分類モデルを作ってみました。
SageMaker AIは、機械学習モデルの学習からデプロイまでをAWS上で行えるサービスです。2024年12月に「Amazon SageMaker」から「Amazon SageMaker AI」に名前が変わりました。
この記事では、次の流れを一通り試します。
- SageMaker AI Studio（JupyterLab）の環境を作る
- MNISTのデータをS3に置く
- PyTorchの学習スクリプトを用意し、学習ジョブとして実行する
- 学習したモデルをリアルタイムエンドポイントにデプロイする
- エンドポイントに画像を送って推論する
- 使ったリソースを削除する
モデルの精度を上げることは目的にしていません。SageMaker AIで「学習 → デプロイ → 推論」を一通り動かしてみることが目的です。
全体像
環境
| 項目 | 内容 |
|---|---|
| リージョン | 米国東部（バージニア北部）us-east-1 |
| SageMaker Python SDK | 3.29.0 |
| 学習コンテナ | PyTorch 2.8 |
| 推論コンテナ | PyTorch 2.6（py312、CPU）※ModelBuilderが自動で選択 |
| ノートブック | SageMaker AI Studio JupyterLab（ml.t3.medium） |
| 学習ジョブ | ml.m5.xlarge × 1（CPU） |
| エンドポイント | ml.m5.large × 1（CPU） |
手順1：SageMaker AIのドメインを作る
SageMaker AI Studioを使うには、まず「ドメイン」を作ります。ドメインは、Studioのユーザーや設定、保存領域をまとめたものです。
- AWSマネジメントコンソールで「Amazon SageMaker AI」を開きます。
- 左メニューの「ドメイン」から「ドメインを作成」を選びます。
- 「シングルユーザー向けの設定（クイックセットアップ）」を選び、「設定」を押します。
- 数分待つと、ドメインとユーザープロファイル、実行ロール（
AmazonSageMaker-ExecutionRole-...）が自動で作られます。
手順2：JupyterLabを起動する
- 作成したドメインのユーザープロファイルから「起動」→「Studio」を選びます。
- Studioが開いたら、「JupyterLab」→「Create JupyterLab space」を選びます。
- 名前（例:
mnist-handson）を入れ、インスタンスはml.t3.mediumのままで「スペースを実行」を押します。 - 起動したら「開く」で JupyterLab を開き、「Python 3」のノートブックを新規作成します。
以降のコードは、すべてこのノートブックのセルで実行します。
手順3：SDKを準備する
SageMaker Python SDKをv3に更新します。
%pip install -U "sagemaker>=3"
続いて、セッション・実行ロール・保存先のS3バケットを用意します。
from sagemaker.core.helper.session_helper import Session, get_execution_role
session = Session()
role = get_execution_role()
region = session.boto_region_name
bucket = session.default_bucket() # sagemaker-<リージョン>-<アカウントID>
prefix = "mnist-handson"
print(region, role, bucket, sep="\n")
us-east-1
arn:aws:iam::xxxxxxxxxxxx:role/service-role/AmazonSageMaker-ExecutionRole-20261109T074175
sagemaker-us-east-1-xxxxxxxxxxxx
手順4：MNISTのデータをS3に置く
ノートブック上でMNISTをダウンロードして、S3にアップロードします。先に torch と torchvision を入れておきます。
%pip install torch torchvision
from torchvision import datasets
datasets.MNIST(root="data", train=True, download=True)
datasets.MNIST(root="data", train=False, download=True)
Dataset MNIST
Number of datapoints: 10000
Root location: data
Split: Test
data/MNIST/raw/ に、学習用6万枚とテスト用1万枚の画像とラベルが保存されます。どんな画像か、いくつか表示してみます。
import matplotlib.pyplot as plt
test_set = datasets.MNIST(root="data", train=False)
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for ax, i in zip(axes, range(8)):
img, label = test_set[i]
ax.imshow(img, cmap="gray")
ax.set_title(label)
ax.axis("off")
plt.show()
S3にアップロードします。
!aws s3 sync data s3://{bucket}/{prefix}/data --quiet
!aws s3 ls s3://{bucket}/{prefix}/data/MNIST/raw/
2026-11-10 22:13:42 7840016 t10k-images-idx3-ubyte
2026-11-10 22:13:42 1648877 t10k-images-idx3-ubyte.gz
2026-11-10 22:13:42 10008 t10k-labels-idx1-ubyte
2026-11-10 22:13:42 4542 t10k-labels-idx1-ubyte.gz
2026-11-10 22:13:42 47040016 train-images-idx3-ubyte
2026-11-10 22:13:42 9912422 train-images-idx3-ubyte.gz
2026-11-10 22:13:42 60008 train-labels-idx1-ubyte
2026-11-10 22:13:42 28881 train-labels-idx1-ubyte.gz
手順5：学習スクリプトを書く
学習ジョブの中で実行するスクリプトを src/train.py として作ります。%%writefile は保存先のフォルダを作ってくれないので、先に src フォルダを作っておきます。
!mkdir -p src
train.py
import argparse
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
class Net(nn.Module):
"""畳み込み2層 + 全結合2層のシンプルなCNN"""
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3)
self.conv2 = nn.Conv2d(32, 64, 3)
self.fc1 = nn.Linear(64 * 12 * 12, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = (x - 0.1307) / 0.3081
x = F.relu(self.conv1(x))
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
return self.fc2(x)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--epochs", type=int, default=3)
parser.add_argument("--batch-size", type=int, default=64)
parser.add_argument("--lr", type=float, default=1e-3)
# SageMaker AIが環境変数で渡してくる入出力のパス
parser.add_argument("--data-dir", default=os.environ.get("SM_CHANNEL_TRAINING", "data"))
parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "model"))
args, _ = parser.parse_known_args()
transform = transforms.ToTensor()
train_set = datasets.MNIST(args.data_dir, train=True, download=False, transform=transform)
test_set = datasets.MNIST(args.data_dir, train=False, download=False, transform=transform)
train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
test_loader = DataLoader(test_set, batch_size=1000)
model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
for epoch in range(1, args.epochs + 1):
model.train()
for x, y in train_loader:
optimizer.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
optimizer.step()
model.eval()
correct = 0
with torch.no_grad():
for x, y in test_loader:
correct += (model(x).argmax(dim=1) == y).sum().item()
print(f"epoch={epoch} loss={loss.item():.4f} test_accuracy={correct / len(test_set):.4f}")
# TorchScript形式で保存する。推論側でクラス定義がなくても読み込める
os.makedirs(args.model_dir, exist_ok=True)
torch.jit.script(model.eval()).save(os.path.join(args.model_dir, "model.pth"))
if __name__ == "__main__":
main()
データとモデルの置き場所は、SageMaker AIが渡してくる環境変数から受け取るようにしています。
| 環境変数 | 中身 |
|---|---|
SM_CHANNEL_TRAINING |
学習ジョブに渡したデータ（training チャネル）が置かれるディレクトリ。/opt/ml/input/data/training |
SM_MODEL_DIR |
ここに保存したファイルが、学習後に model.tar.gz にまとめられてS3に保存される。/opt/ml/model |
環境変数がない場合は data と model を使うので、ノートブック上でもそのまま動かせます。学習ジョブに投げる前に、1エポックだけ手元で動かして確認しました。
!python src/train.py --epochs 1 --data-dir data --model-dir model_local
epoch=1 loss=0.0033 test_accuracy=0.9831
手順6：学習ジョブを実行する
ModelTrainer を使って、学習ジョブを開始します。
from sagemaker.core import image_uris
from sagemaker.train.configs import Compute, InputData, SourceCode
from sagemaker.train.model_trainer import ModelTrainer
training_image = image_uris.retrieve(
framework="pytorch",
region=region,
version="2.8", # 使えるバージョンは Deep Learning Containers の一覧で確認する
py_version="py312",
instance_type="ml.m5.xlarge",
image_scope="training",
)
model_trainer = ModelTrainer(
training_image=training_image,
role=role,
source_code=SourceCode(source_dir="src", entry_script="train.py"),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
input_data_config=[
InputData(channel_name="training", data_source=f"s3://{bucket}/{prefix}/data"),
],
hyperparameters={"epochs": 3, "batch-size": 64, "lr": 0.001},
base_job_name="mnist-pytorch",
)
model_trainer.train()
実行すると、ノートブックに学習ジョブのログが流れます。コンテナの起動に数分かかってから、学習が始まります。
学習ジョブの様子は、コンソールの「SageMaker AI」→「トレーニング」→「トレーニングジョブ」でも確認できます。ログはCloudWatch Logsに残ります。
JupyterLabスペースのインスタンスと、学習ジョブのインスタンスの違い
JupyterLabスペースで選んだ ml.t3.medium と、コードで指定した ml.m5.xlarge は、別々のマシンです。
| JupyterLabスペースで選んだインスタンス | コードで指定した instance_type |
|
|---|---|---|
| 役割 | ノートブックを開いて、コードを書いて実行する作業用のマシン | 学習ジョブを実際に処理する計算用のマシン |
| 動くもの | JupyterLab本体とカーネル（セルの実行、データの前処理、SDKの呼び出し） | 学習コンテナ（train.py などの学習処理） |
| 起動と停止 | スペースを起動している間は動き続ける | ジョブの開始時に自動で起動し、終わると自動で消える |
| 課金 | スペースを起動している時間の分（止め忘れると課金が続く） | ジョブが動いた秒数の分だけ |
| クォータ名 | Studio JupyterLab Apps running on ml.t3.medium など |
ml.m5.xlarge for training job usage など |
| 使えるタイプ | ml.t3.medium など、軽いタイプも選べる | 学習ジョブに対応したタイプだけ（ml.t3 系は使えない） |
JupyterLabスペース (ml.t3.medium)
│ model_trainer.train() を実行する
│ ＝ CreateTrainingJob API を呼ぶだけ
▼
SageMaker AIが学習用インスタンスを新しく起動 (ml.m5.xlarge)
│ S3からデータを読み込み、学習を実行
│ モデルをS3へ保存
▼
ジョブが終わるとインスタンスは自動で削除される
手順7：学習したモデルを確認する
学習ジョブが保存した model.tar.gz の場所を確認します。
import boto3
sm = boto3.client("sagemaker", region_name=region)
job_name = model_trainer._latest_training_job.training_job_name
model_data = sm.describe_training_job(TrainingJobName=job_name)["ModelArtifacts"]["S3ModelArtifacts"]
print(model_data)
s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-pytorch/mnist-pytorch-2026-11-10-22-14-07-898/output/model.tar.gz
ダウンロードして展開します。
!aws s3 cp {model_data} model.tar.gz
!mkdir -p model_artifacts && tar -xzf model.tar.gz -C model_artifacts
!ls model_artifacts
download: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-pytorch/mnist-pytorch-2026-11-10-22-14-07-898/output/model.tar.gz to ./model.tar.gz
tar: Ignoring unknown extended header keyword 'LIBARCHIVE.creationtime'
model.pth
ノートブック上でモデルを読み込んで、テストデータで動くか確かめます。
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
model = torch.jit.load("model_artifacts/model.pth", map_location="cpu")
test_set = datasets.MNIST("data", train=False, transform=transforms.ToTensor())
correct = 0
with torch.no_grad():
for x, y in DataLoader(test_set, batch_size=1000):
correct += (model(x).argmax(dim=1) == y).sum().item()
print(f"test_accuracy={correct / len(test_set):.4f}")
test_accuracy=0.9881
学習ジョブのログに出ていた、最終エポックの test_accuracy と同じ値になりました。
手順8：エンドポイントにデプロイする
ModelBuilder を使って、学習したモデルをリアルタイムエンドポイントにデプロイします。
推論のときに「モデルをどう読み込み、入力をどう処理するか」は、InferenceSpec を継承したクラスに書きます。
load：エンドポイントの起動時に1回だけ呼ばれ、モデルを読み込みます。invoke：リクエストが来るたびに呼ばれ、推論結果を返します。
import time
from importlib.metadata import version
import torch
from sagemaker.serve.builder.schema_builder import SchemaBuilder
from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.serve.spec.inference_spec import InferenceSpec
class MnistSpec(InferenceSpec):
def load(self, model_dir: str):
import torch
return torch.jit.load(f"{model_dir}/model.pth", map_location="cpu")
def invoke(self, input_object, model):
x = torch.tensor(input_object, dtype=torch.float32)
with torch.no_grad():
return model(x).tolist()
sample_input = [[[[0.0] * 28] * 28]]
sample_output = [[0.0] * 10]
model_builder = ModelBuilder(
inference_spec=MnistSpec(),
model_path="model_artifacts",
schema_builder=SchemaBuilder(sample_input, sample_output),
role_arn=role,
dependencies={
"auto": False,
"custom": [
f"cloudpickle=={version('cloudpickle')}",
f"sagemaker-serve=={version('sagemaker-serve')}",
f"sagemaker-core=={version('sagemaker-core')}",
],
},
)
# 名前が重複しないように、時刻を付ける
suffix = time.strftime("%Y%m%d-%H%M%S")
model_name = f"mnist-pytorch-{suffix}"
endpoint_name = f"mnist-pytorch-{suffix}"
model_builder.build(model_name=model_name)
endpoint = model_builder.deploy(
endpoint_name=endpoint_name,
instance_type="ml.m5.large",
initial_instance_count=1,
)
print(endpoint_name)
dependencies には、推論コンテナに追加でインストールするパッケージを書いています。
ModelBuilder は、MnistSpec のインスタンスを cloudpickle で serve.pkl というファイルに保存して、モデルと一緒にS3へアップロードします。推論コンテナはこのファイルを読み込んで MnistSpec を復元するので、コンテナ側にも次のパッケージが必要です。
| パッケージ | 必要な理由 |
|---|---|
cloudpickle |
serve.pkl を読み込むため |
sagemaker-serve |
MnistSpec の親クラス InferenceSpec（sagemaker.serve）を読み込むため |
sagemaker-core |
sagemaker-serve が使っているため |
推論コンテナ（PyTorchのDeep Learning Container）には、これらは最初から入っていません。dependencies に書いたパッケージは code/requirements.txt に書き出され、コンテナの起動時に pip install されます。
"auto": True（既定）にすると必要なパッケージを自動で検出してくれますが、今回は sagemaker-serve が検出されませんでした。そのため、自動検出は使わずに明示的に指定しています（詳しくは「ハマったところ」に書きました）。
ノートブックと同じバージョンを入れるために、importlib.metadata.version() でバージョンを取ってきています。
エンドポイントが使えるようになるまで、数分かかります。コンソールの「SageMaker AI」→「推論」→「エンドポイント」で、ステータスが「InService」になれば準備完了です。
手順9：エンドポイントで推論する
テストデータの画像を8枚送って、予測させてみます。エンドポイントの呼び出しには、boto3の sagemaker-runtime を使います。
test_set は、手順7で transforms.ToTensor() を付けて作ったものを使います。手順4の test_set（transform なし）のままだと、画像がPILの Image のままなので .tolist() でエラーになります。
import json
import boto3
import matplotlib.pyplot as plt
runtime = boto3.client("sagemaker-runtime", region_name=region)
images = [test_set[i][0] for i in range(8)] # 各画像は [1, 28, 28] のテンソル
labels = [test_set[i][1] for i in range(8)]
payload = [img.tolist() for img in images] # [8, 1, 28, 28]
response = runtime.invoke_endpoint(
EndpointName=endpoint_name,
ContentType="application/json",
Body=json.dumps(payload),
)
scores = json.loads(response["Body"].read())
preds = [max(range(10), key=lambda k: s[k]) for s in scores]
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for ax, img, label, pred in zip(axes, images, labels, preds):
ax.imshow(img.squeeze(), cmap="gray")
ax.set_title(f"label={label} / pred={pred}")
ax.axis("off")
plt.show()
エンドポイントから推論結果が返ってきました。
手順10：後片付け
ここで止めないと料金がかかり続けます。必ず実行してください。
エンドポイントを削除する
エンドポイント本体、エンドポイント設定、モデルの3つを削除します。
config_name = sm.describe_endpoint(EndpointName=endpoint_name)["EndpointConfigName"]
model_names = [
v["ModelName"]
for v in sm.describe_endpoint_config(EndpointConfigName=config_name)["ProductionVariants"]
]
sm.delete_endpoint(EndpointName=endpoint_name)
sm.delete_endpoint_config(EndpointConfigName=config_name)
for name in model_names:
sm.delete_model(ModelName=name)
コンソールの「エンドポイント」の一覧から消えていることを確認します。
S3のデータを削除する
アップロードしたMNISTのデータを削除します。
!aws s3 rm s3://{bucket}/{prefix} --recursive
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/t10k-labels-idx1-ubyte
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/t10k-images-idx3-ubyte.gz
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/t10k-labels-idx1-ubyte.gz
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/t10k-images-idx3-ubyte
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/train-labels-idx1-ubyte
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/train-labels-idx1-ubyte.gz
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/train-images-idx3-ubyte.gz
delete: s3://sagemaker-us-east-1-xxxxxxxxxxxx/mnist-handson/data/MNIST/raw/train-images-idx3-ubyte
同じバケットには、学習ジョブの出力（mnist-pytorch/）と、ModelBuilder がアップロードしたモデル（model-builder/）も残っています。不要であれば、これらも削除します。
!aws s3 rm s3://{bucket}/mnist-pytorch --recursive
!aws s3 rm s3://{bucket}/model-builder --recursive
JupyterLabスペースを止める
Studioの「JupyterLab」の画面で、スペースの「停止」を押します。止めている間は、ノートブックのインスタンス料金はかかりません。
ドメインを削除する（今後使わない場合）
ドメインを残しておくと、ユーザーのファイルを保存しているEFS（Amazon Elastic File System）のストレージ料金がかかります。今後使わない場合は削除します。
- コンソールでドメインを開き、「スペース管理」からスペースを削除します。
- 「ユーザープロファイル」を削除します。
- ドメインを削除します。
- 削除後、EFSのコンソールで、ドメイン用に作られたファイルシステムが残っていないか確認します。
ハマったところ
学習ジョブのクォータが0になっていた
モデル学習時に次のエラーが出ました。
ResourceLimitExceeded: An error occurred (ResourceLimitExceeded) when calling the CreateTrainingJob operation: The account-level service limit 'ml.m5.xlarge for training job usage' is 0 Instances, with current utilization of 0 Instances and a request delta of 1 Instances. Please use AWS Service Quotas to request an increase for this quota. If AWS Service Quotas is not available, contact AWS support to request an increase for this quota.
SageMaker AIを初めて使うアカウントだったため、学習ジョブ用のクォータが0になっていました。次の手順で引き上げを申請しました。
- 「Service Quotas」→「AWSのサービス」→「Amazon SageMaker」を開く
- 「ml.m5.xlarge for training job usage」を選ぶ
- 「クォータの引き上げをリクエスト」で
1を入力する（並列で動かす場合は、その数を入力する）
1時間くらいで承認されました。
クォータは用途ごとに分かれていて、エンドポイント用は「ml.m5.large for endpoint usage」のように別の項目になっています。
デプロイで推論コンテナが起動しなかった
一番時間がかかったのがデプロイです。最初は dependencies を指定せずに build() を実行して、次のエラーになりました。
CalledProcessError: Command '['/opt/conda/bin/python', '-c', "from sagemaker.serve.detector.pickle_dependencies import get_requirements_for_pkl_file;get_requirements_for_pkl_file('/home/sagemaker-user/model_artifacts/code/serve.pkl', '/home/sagemaker-user/model_artifacts/code/requirements.txt')"]' returned non-zero exit status 1.
ここから、次の順で問題が出ました。
| 出た症状 | 原因（AI解説） | 対応（AI解説） |
|---|---|---|
build() で上の CalledProcessError |
依存パッケージを自動検出する処理は、環境変数を引き継がない別プロセスで動く。そのためリージョンが分からず、NoRegionError で落ちていた |
!aws configure set region us-east-1 で、設定ファイル（~/.aws/config）にリージョンを書いた |
推論がタイムアウトし、コンテナのログに No module named 'cloudpickle' |
cloudpickle がコンテナに入っていなかった。さらに、同じモデル名（mnist-pytorch）で作り直していたため、Using already existing model となり、古いモデルが使われ続けていた |
dependencies で cloudpickle を指定し、モデル名とエンドポイント名に時刻を付けた |
推論がタイムアウトし、コンテナのログに No module named 'sagemaker.serve' |
自動検出で作られた requirements.txt に、sagemaker-core と sagemaker-train は入っていたが、sagemaker-serve が入っていなかった |
dependencies で sagemaker-serve を明示的に指定した |
推論がタイムアウトしたときは、ノートブック側にはタイムアウトのエラーしか出ません。原因は、CloudWatch Logsの /aws/sagemaker/Endpoints/<エンドポイント名> に出ているコンテナのログで確認しました。
また、エンドポイントの作成に失敗したあと、同じ名前でもう一度デプロイすると、次のエラーになります。
ClientError: An error occurred (ValidationException) when calling the CreateEndpointConfig operation: Cannot create already existing endpoint configuration
失敗したエンドポイントでも、エンドポイント設定は残っています。作り直すときは、エンドポイントとエンドポイント設定を削除するか、別の名前を使います。
おわりに
SageMaker AIで、学習ジョブからエンドポイントでの推論まで一通り動かしました。
学習ジョブは終わると自動でインスタンスが消えるので、ノートブックは ml.t3.medium のままで進められました。デプロイでは何度かつまずきましたが、CloudWatch Logsでコンテナのログを見ることで原因を確認できました。