🕶️
AWS Sagemaker Studioで独自Containerを使う-前編-
はじめに
今日、クラウドでDeep Learning Modelのトレーニングを行うことは必須技術となってきました。
AWSを始めGCPやAzureでも似たようなサービスを提供していますが、S3やEC2などのAWSのサービスをしている場合、やはりトレーニング自体もAWS内で行うことがスムーズです。
今回は、独自のモデルをDockerを使用して学習する方法を説明します。
セットアップ
AWSコンソールからSageMakerを選択し、"今すぐ始める"ボタンでセットアップを行います。
ここではSagemaker Domainの設定だけで十分です。
その後、コントロールパネルからstudioを選択し、Sagemaker Studioを起動します。
ユーザーを追加でユーザーを作成した後、"アプリケーションを起動"から Sagemaker Studioを起動します。
トレーニングコード
さて、ここでトレーニングスクリプトを準備します。PyTorchやTensorflowなど好きなフレームワークを使いトレーニングスクリプトを書きます。
その前にここで作成するファイルをまとめるために、ローカルマシンに新規フォルダを作成し、下記のように空のファイルを作成してみましょう。
MinstTraining/
┣ code/train.py
┣ Dockerfile
こでは単純な画像分類を例にします。
下記の内容をtrain.pyにコピーします。
train
from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets, transforms
import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler(sys.stdout))
def build_model(
weights='IMAGENET1K_V2',
fine_tune=False,
num_classes=10
):
if weights:
print(f"[INFO]: Loading {weights} pre-trained weights")
else:
print('[INFO]: Not loading pre-trained weights')
model = models.efficientnet_b1(weights=weights)
if fine_tune:
print('[INFO]: Fine-tuning all layers...')
for params in model.parameters():
params.requires_grad = True
elif not fine_tune:
print('[INFO]: Freezing hidden layers...')
for params in model.parameters():
params.requires_grad = False
# Change the final classification head.
model.classifier[1] = nn.Linear(in_features=1280, out_features=num_classes)
return model
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
train_loss = loss.item()
logger.info(f"Epoch: {epoch}, Train Loss: {train_loss:.6f},\n")
if args.dry_run:
break
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
test_acc = 100. * correct / len(test_loader.dataset)
logger.info(f"Test Loss: {test_loss:.6f}, Test Acc: {test_acc:.1f};\n")
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument("--work-dir", type=str, default="/opt/ml/model",
help="the dir to save logs and models")
parser.add_argument('--data-dir', type=str,
default="/opt/ml/input/data/train", help="Data dir")
parser.add_argument('--batch-size', type=int, default=32, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--epochs', type=int, default=20, metavar='N',
help='number of epochs to train (default: 20)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--dry-run', action='store_true', default=False,
help='quickly check a single pass')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
args = parser.parse_args()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
train_dataset = ImageFolder(f'{args.data_dir}/train', transform)
val_dataset = ImageFolder(f'{args.data_dir}/val', transform)
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
sampler=sampler,
shuffle=shuffle,
num_workers=4,
pin_memory=True,
drop_last=True
)
val_loader = DataLoader(
val_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True
)
num_classes = len(train_dataset.classes)
model = build_model(num_classes=num_classes)
model = model.to('cuda:0')
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(model, device, test_loader)
scheduler.step()
torch.save(model.state_dict(), f"{args.work_dir}/model.pt")
if __name__ == '__main__':
main()
Dockerfile
上記のコードを実行するために、Dokerfileを作成します。
# https://hub.docker.com/r/pytorch/pytorch
FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-devel
RUN apt-get update
RUN pip install --upgrade pip
# /opt/ml and all subdirectories are utilized by SageMaker, we use the /code subdirectory to store our user code.
COPY /code /opt/ml/code
# Install sagemaker-training module
RUN pip install --no-cache --upgrade \
sagemaker-training
# Defines train.py as script entrypoint
ENV SAGEMAKER_PROGRAM train.py
PyTorchのコンテナをPullして使用します。https://hub.docker.com/r/pytorch/pytorch/tags
また、pipでsagemaker-trainingをインストールし、
ENV SAGEMAKER_PROGRAM train.py
これでエントリーポイントとなるpythonなどのファイルを指定します。これが学習時の実行ファイルとなります。
この後、Conatinerのビルドと学習の実行を行いますが、IAMの設定など少し長くなるので、ここで一旦区切ります。
次回、AWS Sagemaker Studioで独自Containerを使う-中編-へ続きます。
修正・追記
2022/08/30 codeの一部を修正
Discussion