AWS Sagemakerで自作モデルを使用し学習・モデル作成・デプロイまで実装してみた

機械学習エンジニアであるもののSageMakerは試したことがなかったので、この機会に実装してみた。

学習用データは、Kaggleの以下データを使用した。

Cervical Cancer Risk Classification:
https://www.kaggle.com/loveall/cervical-cancer-risk-classification

フレームワークは、個人的にPytorch/Chainerが得意なので、Pytorchでネットワークを組んでみた。

AWSでJupyterインスタンスがあるが、今月は金欠なのでローカルで開発することにした。。
まずは、SageMakerに対するポリシーを付与したIAMロールを作成し、実行アカウント(IAMユーザ)・SageMaker双方への信頼関係を付与する。
当ロールに、以下のポリシーをアタッチする。
・ IAMReadOnlyAccess
・ AmazonSageMakerFullAccess

```
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "",
      "Effect": "Allow",
      "Principal": {
        "AWS": "arn:aws:iam::xxxxx:root",
        "Service": "sagemaker.amazonaws.com"
      },
      "Action": "sts:AssumeRole"
    }
  ]
}
```

作成したロールの実行ポリシーをIAMユーザへ割り当てる。

```
{
    "Version": "2012-10-17",
    "Statement": {
        "Effect": "Allow",
        "Action": "sts:AssumeRole",
        "Resource": "arn:aws:iam::xxxxx:role/sagemaker-role"
    }
}
```


ローカルの.aws/configに上記で作成したロールを付与するよう設定する。

```
[default]
output = json
region = ap-northeast-1
role_arn = arn:aws:iam::xxxxx:role/sagemaker-role
source_profile = default
```

これで、ローカルからSageMakerへのアクセス権限の設定は完了したので、後はいつも通りJupyterで学習用ネットワークを実装すればよい。


```Python
import numpy as np
import subprocess
import sys 
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim
import argparse
import sagemaker_containers
import os

class Model(nn.Module):
  def __init__(self, input_dim, output_dim):
    super(Model, self).__init__()
    self.sequential = nn.Sequential(
      nn.Linear(input_dim, 40),
      nn.ReLU(),
      nn.Dropout(0.3),
      nn.BatchNorm1d(40),
      nn.Linear(40, 20),
      nn.ReLU(),
      nn.Dropout(0.3),
      nn.BatchNorm1d(20),
      nn.Linear(20,10),
      nn.ReLU(),
      nn.Dropout(0.3),
      nn.BatchNorm1d(10),
      nn.Linear(10,output_dim),
      )

  def forward(self, x):
    x = self.sequential(x)

    return F.log_softmax(x, dim=1)

def preprocess(args):
    subprocess.call([sys.executable, "-m", "pip", "install", "scikit-learn"]) 
    subprocess.call([sys.executable, "-m", "pip", "install", "pandas==0.23.0"]) 
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import MinMaxScaler  
    import pandas as pd
    print(pd.__version__)
    df = pd.read_csv(args.data_dir + '/kag_risk_factors_cervical_cancer.csv') 
    df = df.convert_objects(convert_numeric=True)
    dum_list = ['Smokes', 'IUD', 'STDs', 'Hormonal Contraceptives','Dx:Cancer','Dx:CIN','Dx:HPV','Dx','Hinselmann','Citology','Schiller']
    for column in df.columns.values:
        df[column] = df[column].fillna(df.median()[column])
        if column in dum_list and column != 'Biopsy':
            dummy_df = pd.get_dummies(df[column])
            df = pd.concat([df, dummy_df], axis=1)
            df = df.drop(column, axis=1)

    X_df_data = df.drop('Biopsy', axis=1)
    y_df_data = df['Biopsy']

    X_train, X_test = train_test_split(np.array(X_df_data), test_size=0.2, random_state=0)
    y_train, y_test = train_test_split(np.array(y_df_data), test_size=0.2, random_state=0)

    ms = MinMaxScaler(feature_range=(0, 1))
    X_train, X_test = ms.fit_transform(X_train), ms.fit_transform(X_test)

    X_train_tensor, X_test_tensor, y_train_tensor, y_test_tensor = torch.from_numpy(X_train).float(), torch.from_numpy(X_test).float(), torch.from_numpy(y_train).long(), torch.from_numpy(y_test).long()

    train_data = TensorDataset(X_train_tensor, y_train_tensor)
    train_loader = DataLoader(train_data, batch_size=
                              args.train_batch_size, shuffle=True)

    test_data = TensorDataset(X_test_tensor, y_test_tensor)
    test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=True)

    return train_loader, X_train_tensor, test_loader, y_test, y_test_tensor

def train(args):
    train_loader, X_train_tensor, test_loader, y_test, y_test_tensor = preprocess(args)
    model = Model(input_dim=X_train_tensor.shape[1], output_dim=args.num_classes)

    criterion = nn.CrossEntropyLoss()
    optim = torch.optim.Adam(model.parameters(), lr=args.lr)

    for epoch in range(args.epochs):
      total_loss = 0
      for X_train, y_train in train_loader:
        X_train, y_train = Variable(X_train), Variable(y_train)
        optim.zero_grad()
        output = model(X_train)
        loss = criterion(output, y_train)
        loss.backward()
        optim.step()

        total_loss += loss.item()

      if (epoch + 1) % 10 == 0:
        print(epoch + 1, total_loss)

    return _save_model(model, args.model_dir), model, test_loader, y_test, y_test_tensor

def accuracy(model, test_loader, y_test, y_test_tensor):
    with torch.no_grad():
      tp_cnt, tn_cnt, correct_cnt = 0, 0, 0
      for inputs, labels in test_loader:
        inputs, labels = Variable(inputs), Variable(labels)

        _, predict = torch.max(model(inputs), 1)
        for val, label in zip(predict, labels):
          if val.item() == label:
            correct_cnt += 1
            if val.item() == 1:
              tp_cnt += 1
            elif val.item() == 0:
              tn_cnt += 1
    total_acc = correct_cnt / y_test_tensor.size()[0]
    tp_acc = tp_cnt / np.count_nonzero(y_test == 1)
    tn_cnt = tn_cnt / np.count_nonzero(y_test == 0)
    print('Total accuracy : ' + str(round(total_acc, 2)) + '\nTP : ' + str(round(tp_acc, 2)) + '\nTN : ' + str(round(tn_cnt, 2)))

def _save_model(model, model_dir):
    path = os.path.join(model_dir, 'model.pth')
    torch.save(model.cpu().state_dict(), path)

def model_fn(model_dir):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model(input_dim=46, output_dim=2)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)

    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
        model.load_state_dict(torch.load(f))
    return model.to(device)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=100, metavar='E',
                    help='number of total epochs to run (default: 100)')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                    help='initial learning rate (default: 0.001)')
    parser.add_argument('--train-batch-size', type=int, default=200, metavar='E',
                help='number of train data batch size (default: 200)')
    parser.add_argument('--test-batch-size', type=int, default=200, metavar='E',
                help='number of test data batch size (default: 10)')
    parser.add_argument('--num-classes', type=int, default=2, metavar='E',
                help='number of classes (default: 2)')

    env = sagemaker_containers.training_env()
    parser.add_argument('--hosts', type=list, default=env.hosts)
    parser.add_argument('--current-host', type=str, default=env.current_host)
    parser.add_argument('--model-dir', type=str, default=env.model_dir)
    parser.add_argument('--data-dir', type=str, default=env.channel_input_dirs.get('training'))
    parser.add_argument('--num-gpus', type=int, default=env.num_gpus)
    
    _, model, test_loader, y_test, y_test_tensor = train(parser.parse_args())
    accuracy(model, test_loader, y_test, y_test_tensor)

```

SageMakerのPytorchフレームワークから上記のスクリプトを使用した学習・エンドポイントへのデプロイを実施するよう実装する。

```Python
import subprocess
from sagemaker import Session, get_execution_role
from sagemaker.pytorch import PyTorch
import sagemaker

class SageMaker:
    def __init__(self):
        self.sagemaker_session = sagemaker.Session()
        self.bucket = '*****' #S3バケット名
        self.role = sagemaker.get_execution_role()
        self.instance_type = '*****' #デプロイ先インスタンスタイプ
    
    def upload_data(self):
        inputs = self.sagemaker_session.upload_data(path='./data/', bucket=self.bucket, key_prefix='data/cancer')
        
        return inputs
    
    def __estimator_fit(self, inputs):
        estimator = PyTorch(entry_point='./kag_risk_train_test.py',
                            role=self.role,
                            framework_version='1.1.0',
                            train_instance_count=1,
                            train_instance_type=self.instance_type)

        estimator.fit(
            inputs=inputs,
            )
        
        return estimator
        
    def estimator_predict(self, inputs):
        estimator = self.__estimator_fit(inputs)
        predict = estimator.deploy(initial_instance_count=1,
                                             instance_type=self.instance_type)
        
        return predict

```

実行結果:
全体的な精度は良いがTP値が少し低いので、もうちょい改善の余地はあるかも。。

```
10 2.986310064792633
20 2.434782922267914
30 1.926128625869751
40 1.4562886655330658
50 1.0403850376605988
60 0.7982784509658813
70 0.6503457129001617
80 0.6562095209956169
90 0.4315394461154938
100 0.46829692274332047

Total accuracy : 0.96
TP : 0.78
TN : 0.97
2019-10-25 16:57:42,283 sagemaker-containers INFO     Reporting training SUCCESS
```

学習・デプロイ(エンドポイント作成)が成功すると、以下の通りトレーニングジョブステータスがCompleted、エンドポイントのステータスがInServiceとなる。

![スクリーンショット 2019-10-26 2.06.00（2）.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/7cdcdc57-3592-a846-f7c7-187e74643589.png)

![スクリーンショット 2019-10-26 2.25.35（2）.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/6c1c3736-e9e8-54b3-449c-fca0450071fd.png)