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

## 記事一覧

https://qiita.com/kazuneet/items/795e561efce8c874d115

https://qiita.com/kazuneet/items/294357a2214ac733f6f1

https://qiita.com/kazuneet/items/caa887826e2c33309e0b

https://qiita.com/kazuneet/items/62edfa445e526eb322d6

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

[前回の記事](https://qiita.com/kazuneet/items/62edfa445e526eb322d6)では TensorFlow のトレーニングコードを SageMaker Training に向けてコードを書き換えましたが、今回は PyTorch 版です。

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

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

https://github.com/kazuhitogo/qiita-sagemaker-training

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

https://qiita.com/kazuneet/items/62edfa445e526eb322d6

## トレーニングコード

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

```python:./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 でそのまま動きます。

```text:実行結果
(略)
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` に用意して以下のコードを実行します。

```python:./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()
```

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

```text:実行結果
(略)
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 に転送してくれます。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E5%87%A6%E7%90%86%E7%B5%90%E6%9E%9C%E3%82%92%E4%BF%9D%E5%AD%98%E3%81%99%E3%82%8B%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E3%81%A7%E8%A8%80%E3%81%86%E3%81%A8%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E4%BF%9D%E5%AD%98

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

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

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

```diff:./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 で動かしてみましょう。

```python:./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)
```

```text:実行結果
(略)
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に配置したデータからトレーニングインスタンスに転送）やることで、トレーニングジョブ呼び出し時に使用するデータを決定し、トレーニングコードは与えられたデータを読み込むだけの実装に変更してみましょう。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E7%94%A8%E6%84%8F%E3%81%97%E3%81%9F%E3%83%87%E3%83%BC%E3%82%BF%E3%82%92-sagemaker-training-%E5%86%85%E3%81%A7%E4%BD%BF%E3%81%86%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E3%81%A7%E8%A8%80%E3%81%86%E3%81%A8%E5%AD%A6%E7%BF%92%E3%83%87%E3%83%BC%E3%82%BF%E3%81%AE%E5%88%A9%E7%94%A8

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

```python:./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` の値が示すディレクトリに転送されるようにジョブを設定しますので、それぞれのディレクトリに各ファイルが存在する前提でコードを記載します。

```diff:./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)

```

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

```python:./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 で示されたデータが環境変数の値で示されるディレクトリに保存されます。

https://qiita.com/kazuneet/items/795e561efce8c874d115#%E8%A4%87%E6%95%B0%E3%81%AE-prefix-%E3%81%AB%E3%81%82%E3%82%8B%E3%83%87%E3%83%BC%E3%82%BF%E3%82%92%E4%BD%BF%E3%81%86-training-job

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

```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)
```

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

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

```diff:./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()を使って定義されていない場合のデフォルト値をセットしています。

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


```python:./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,
})
```

```text:実行結果
(略)
{'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
```
無事完了しました。ハイパーパラメータが効いているか、モデルをロードしてフィルター数だけ確認しておきましょう。
```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, 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)
```
```text:実行結果
(略)
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