LoginSignup
3
4

More than 1 year has passed since last update.

@kazuneet(和仁 呉)

エンジニア目線で始める Amazon SageMaker Training ④-2 PyTorch のコードを SageMaker Training 用へ書き換える

Last updated at Posted at 2022-03-11

記事一覧

既存 PyTorch のトレーニングコードを SageMaker Training 用に書き換える

前回の記事では TensorFlow のトレーニングコードを SageMaker Training に向けてコードを書き換えましたが、今回は PyTorch 版です。

PyTorch のトレーニングコードを SageMaker Training 用に変更します。

※不本意ながらついに機械学習のコードを使います
※例によってコードはすべて GitHub にあるので、Cloneしてお使いください。

※この記事では PyTorch を使っています。TensorFlow(Keras)をお使いの方は以下記事をご参照ください。

トレーニングコード

PyTorch でこんなコードを用意しました。(TensorFlow の時と同等です）

./4_2_rewriting_traing_code_for_sagemaker_pytorch
import torch
import torchvision

# データロード
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
validset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
validloader = torch.utils.data.DataLoader(validset, batch_size=4,shuffle=False, num_workers=2)

#モデル定義
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 16, kernel_size=(3,3), stride=1, padding=(1,1)),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(16*32*32,10),
    torch.nn.Softmax(dim=1)
)

# 学習
def exec_epoch(loader, model, train_flg, optimizer, criterion):
    total_loss = 0.0
    correct = 0
    count = 0
    for i, data in enumerate(loader, 0):
        inputs, labels = data
        if train_flg:
            inputs, labels = torch.autograd.Variable(inputs), torch.autograd.Variable(labels)
            optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        if train_flg:
            loss.backward()
            optimizer.step()
        total_loss += loss.item()
        pred_y = outputs.argmax(axis=1)
        correct += sum(labels==pred_y)
        count += len(labels)
    total_loss /= (i+1)
    total_acc = 100 * correct / count
    if train_flg:
        print(f'train_loss: {total_loss:.3f} train_acc: {total_acc:.3f}%',end=' ')
    else:
        print(f'valid_loss: {total_loss:.3f} valid_acc: {total_acc:.3f}%')
    return model

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.00001)
for epoch in range(2):
    print(f'epoch: {epoch+1}',end=' ')
    model = exec_epoch(trainloader, model, True, optimizer, criterion)
    exec_epoch(validloader, model, False, optimizer, criterion)

# モデル保存
torch.save(model.state_dict(),'1.pth')

データはみんな大好き MNI… ではなく意味もなく CIFER-10 を使ってみました。
CIFER-10 のデータセットを PyTorch から呼び出しています。

モデルはフィルター数 16 の畳み込みを一回だけ入れ、活性化関数に ReLU を指定して、最後に全結合層を入れて活性化関数に Softmax を入れて分類問題の学習をしているだけの、ものすごく簡単な構造です。

最後にモデルを pth 形式で保存しています。

このトレーニングコードは Jupyter でそのまま動きます。

実行結果
(略)
epoch: 1 train_loss: 2.207 train_acc: 27.622% valid_loss: 2.140 valid_acc: 35.550%
epoch: 2 train_loss: 2.120 train_acc: 36.440% valid_loss: 2.101 valid_acc: 37.770%

1. SageMaker Training で動かす

実は特にコードを変更しなくても動きます。
全く同じコードを、src/4-2-1/train.py に用意して以下のコードを実行します。

./4_2_rewriting_traing_code_for_sagemaker_pytorch
import sagemaker
from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point='./src/4-2-1/train.py',
    py_version='py38', 
    framework_version='1.10.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit()

実行すると出力が若干乱れますが、正常終了します。

実行結果
(略)
train_loss: 2.198 train_acc: 28.288%
valid_loss: 2.139 valid_acc: 35.220%
epoch: 2
train_loss: 2.126 train_acc: 35.202%
valid_loss: 2.110 valid_acc: 36.450%
2022-03-09 12:05:16,062 sagemaker-training-toolkit INFO     Reporting training SUCCESS
(略)

というわけで、SageMaker Training 用に特に書き換える必要もなく処理は動くのですが、SageMaker Training は揮発する(処理が終わるとインスタンスが消失する）ので、せっかくトレーニングしたモデルも揮発してしまいます。しかし、SageMaker Training では、①の記事の通り、トレーニングの成果物を指定ディレクトリに保存すれば Amazon S3 に転送してくれます。

そこを実装してみましょう。

2.トレーニングが完了したらモデルを S3 に保存する

早速トレーニングコードを修正してみましょう。

./src/4-2-2/train.pyと./src/4-2-1/train.pyのdiff
(略)

+ import os

(略)


+ model_dir = os.environ.get('SM_MODEL_DIR')
- torch.save(model.state_dict(),'1.pth')
+ torch.save(model.state_dict(),os.path.join(model_dir,'1.pth'))

①の記事の通り、SM_MODEL_DIR という環境変数で指定されているディレクトリに保存すると、処理完了時に Amazon S3 にデータが転送されるため、環境変数の取得と保存先の変更を行いました。

このトレーニングコードを SageMaker で動かしてみましょう。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
import sagemaker
from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point='./src/4-2-2/train.py',
    py_version='py38', 
    framework_version='1.10.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit()

無事トレーニングが終了しますが、Amazon S3 にモデルが転送されたか、モデルを読み込んで確認してみましょう。

python./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import torch
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 16, kernel_size=(3,3), stride=1, padding=(1,1)),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(16*32*32,10),
    torch.nn.Softmax(dim=1)
)
model.load_state_dict(torch.load('1.pth'))
print(model)
実行結果
(略)
Sequential(
  (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU()
  (2): Flatten(start_dim=1, end_dim=-1)
  (3): Linear(in_features=16384, out_features=10, bias=True)
  (4): Softmax(dim=1)
)

無事トレーニング完了後にモデル(重みとバイアス)が S3 に転送され、使えることが確認できました。

3. トレーニングデータをコードの外に出す

機械学習のトレーニングではデータが変わることが多いですが、固定データの呼び出しがコードの中に埋め込まれていると毎回コードを変更する必要が出てきてしまうため、使い勝手が良くないです。
以下記事で紹介した通り、学習データを外から与えて(S3に配置したデータからトレーニングインスタンスに転送）やることで、トレーニングジョブ呼び出し時に使用するデータを決定し、トレーニングコードは与えられたデータを読み込むだけの実装に変更してみましょう。

まずはデータを S3 に保存します。今回は PyTorch テンソル(pt形式)で保存します。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
import torchvision
import numpy as np
import torch
import sagemaker
import os
train_dir = './train'
valid_dir = './valid'
!mkdir -p {train_dir}
!mkdir -p {valid_dir}
train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=len(train_data))
valid_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
valid_data_loader = torch.utils.data.DataLoader(valid_data, batch_size=len(valid_data))
train_data_loaded = next(iter(train_data_loader))
torch.save(train_data_loaded, os.path.join(train_dir, 'train.pt'))
valid_data_loaded = next(iter(valid_data_loader))
torch.save(valid_data_loaded, os.path.join(valid_dir, 'valid.pt'))
train_s3_uri = sagemaker.session.Session().upload_data(path=train_dir, bucket=sagemaker.session.Session().default_bucket(), key_prefix='training/4-2/train')
valid_s3_uri = sagemaker.session.Session().upload_data(path=valid_dir, bucket=sagemaker.session.Session().default_bucket(), key_prefix='training/4-2/valid')

これで train_s3_uri と valid_s3_uri が指す S3 の URI に pt ファイル(train.pt, valid.pt) をアップロードすることができました。

次にトレーニングコードを書き換えます。
後述しますが、train.pt は 環境変数 SM_CHANNEL_TRAIN の値が示すディレクトリに、valid.pt は 環境変数 SM_CHANNEL_VALID の値が示すディレクトリに転送されるようにジョブを設定しますので、それぞれのディレクトリに各ファイルが存在する前提でコードを記載します。

./src/4-2-3/train.pyと./src/4-2-2/train.pyのdiff
(略)

+ train_dir = os.environ.get('SM_CHANNEL_TRAIN')
+ valid_dir = os.environ.get('SM_CHANNEL_VALID')
- train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
+ train_data = torch.utils.data.TensorDataset(*torch.load(os.path.join(train_dir, 'train.pt')))
train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=4, shuffle=True, num_workers=2)
- valid_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]))
+ valid_data = torch.utils.data.TensorDataset(*torch.load(os.path.join(valid_dir, 'valid.pt')))
valid_data_loader = torch.utils.data.DataLoader(valid_data, batch_size=4,shuffle=False, num_workers=2)

これで書き換え終了です。トレーニングジョブを起動しましょう。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
import sagemaker
from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point='./src/4-2-3/train.py',
    py_version='py38', 
    framework_version='1.10.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit({
    'train':train_s3_uri,
    'valid':valid_s3_uri,
})

先程、SM_CHANNEL_TRAIN や SM_CHANNEL_VALID の値が示すディレクトリにファイルが転送されるように設定すると記載しました。
①の記事の通り fit() メソッドの引数に S3 の URI を指定するとトレーニングインスタンスにデータが転送されますが、転送先の指定を行うには引数を辞書形式で渡します。
辞書のキーが環境変数に設定され SM_CHANNEL_{キー} という形で引き渡されます。上のコードの例であれば、train というキーで指定された場合、SM_CHANNEL_TRAIN という環境変数の変数名に、valid というキーで指定された場合はSM_CHANNEL_VALIDという環境変数の変数名で受け渡され、辞書の値で指定された S3 の URI で示されたデータが環境変数の値で示されるディレクトリに保存されます。

処理が無事完了したら、念の為モデルを確認しましょう。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import torch
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 16, kernel_size=(3,3), stride=1, padding=(1,1)),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(16*32*32,10),
    torch.nn.Softmax(dim=1)
)
model.load_state_dict(torch.load('1.pth'))
print(model)
実行結果
(略)
Sequential(
  (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU()
  (2): Flatten(start_dim=1, end_dim=-1)
  (3): Linear(in_features=16384, out_features=10, bias=True)
  (4): Softmax(dim=1)
)

しっかりモデルが出来上がったことがわかります。

4. ハイパーパラメータを外から渡す

だいぶ SageMaker っぽい使い方になってきましたが、もう少し頑張ってみましょう。
モデルを試行錯誤する上でハイパーパラメータの値は変更しやすいと嬉しいです。そのハイパーパラメータ変更のたびに都度コードを変更するのは効率が良くないので、トレーニングジョブ起動時に設定するようにしましょう。また、設定しなくても動くようにデフォルト値も設定しておきましょう。
このケースでは、畳み込みのフィルター数と、バッチサイズ、エポック数と、学習率を外部から与えてみます。

トレーニングコードをこのように修正しました。

./src/4-2-4/train.pyと./src/4-2-3/train.pyのdiff
(略)

+ # ハイパーパラメータ取得
+ hps = json.loads(os.environ.get('SM_HPS'))
+ hps.setdefault('batch_size', 4)
+ hps.setdefault('epochs', 2)
+ hps.setdefault('filters', 16)
+ hps.setdefault('learning_rate', 0.00001)
+ print(hps)
(略)



- train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=4, shuffle=True, num_workers=2)
+ train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=hps['batch_size'], shuffle=True, num_workers=2)
(略)

- valid_data_loader = torch.utils.data.DataLoader(valid_data, batch_size=4, shuffle=True, num_workers=2)
+ valid_data_loader = torch.utils.data.DataLoader(valid_data, batch_size=hps['batch_size'], shuffle=True, num_workers=2)
(略)



-     torch.nn.Conv2d(3, 16, kernel_size=(3,3), stride=1, padding=(1,1)),
+     torch.nn.Conv2d(3, hps['filters'], kernel_size=(3,3), stride=1, padding=(1,1)),
(略)



-     torch.nn.Linear(16*32*32,10),
+     torch.nn.Linear(hps['filters']*32*32,10),
(略)



- optimizer = torch.optim.Adam(params=model.parameters(), lr=0.00001)
- for epoch in range(2):
+ optimizer = torch.optim.Adam(params=model.parameters(), lr=hps['learning_rate'])
+ for epoch in range(hps['epochs']):

ハイパーパラメータを取得するやり方はコマンドライン引数から取得する方法(argparse を使ってパースする)と環境変数から取得する方法がありますが、
個人の好みで環境変数から取得してみました。併せてsetdefault()を使って定義されていない場合のデフォルト値をセットしています。

さてトレーニングコードを実行してみます。ハイパーパラメータが外から与えられているか確認するため、デフォルト値とは別のハイパーパラメータを与えます。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
import sagemaker
from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point='./src/4-2-4/train.py',
    py_version='py38', 
    framework_version='1.10.0',
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role(),
    hyperparameters={
        'filters':8,
        'epochs':1,
        'batch-size':'16',
        'learning_rate' : 0.001
    }
)
estimator.fit({
    'train':train_s3_uri,
    'valid':valid_s3_uri,
})
実行結果
(略)
{'batch_size': 16, 'epochs': 1, 'filters': 8, 'learning_rate': 0.001, 'model_dir': 's3://sagemaker-{REGION}-{ACCOUNT_ID}/tensorflow-training-{YYYY-MM-DD-HH-MI-SS-mmm}/model'}
(略)
train_loss: 2.099 train_acc: 35.726%
valid_loss: 2.025 valid_acc: 43.400%
2022-03-11 03:08:04,725 sagemaker-training-toolkit INFO     Reporting training SUCCESS

2022-03-11 03:08:32 Completed - Training job completed

無事完了しました。ハイパーパラメータが効いているか、モデルをロードしてフィルター数だけ確認しておきましょう。

./4_2_rewriting_traing_code_for_sagemaker_pytorch.ipynb
# モデルの URI を取得
model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']
# モデルをローカルにコピーして解凍
!aws s3 cp {model_uri} ./
!tar zxvf model.tar.gz
# モデルの読み込み
import torch
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 8, kernel_size=(3,3), stride=1, padding=(1,1)),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(8*32*32,10),
    torch.nn.Softmax(dim=1)
)
model.load_state_dict(torch.load('1.pth'))
print(model)
実行結果
(略)
Sequential(
  (0): Conv2d(3, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU()
  (2): Flatten(start_dim=1, end_dim=-1)
  (3): Linear(in_features=8192, out_features=10, bias=True)
  (4): Softmax(dim=1)
)

無事 Conv2d(3, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) とフィルター数が8でパラメータをロードできることを確認できました。

終わりに

PyTorch の機械学習コードを SageMaker で動かすのに適したコードに書き換えました。
ここまでの機能を利用することで、SageMaker のハイパーパラメータチューニングの機能も簡単に利用できるようになるので、ぜひこの使い方を覚えましょう。

https://aws.amazon.com/jp/blogs/news/sagemaker-automatic-model-tuning/
https://sagemaker.readthedocs.io/en/stable/api/training/tuner.html

3
4
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
3
4