#!/bin/bash
sudo -u ec2-user -i <<'EOF'

TARGET_ENV=python3
PROJECT_NAME=myProject

# install package
# This will affect only the Jupyter kernel called "conda_python3".
source activate ${TARGET_ENV}
# Replace myPackage with the name of the package you want to install.
pip install myPackage
# You can also perform "conda install" here as well.
source deactivate

# mkdir
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/configs
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/data
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/features
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/logs
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/models
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/notebook
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/scripts
mkdir -p /home/ec2-user/SageMaker/${PROJECT_NAME}/utils

EOF

data
├── raw          <- The original, immutable data dump.
│   └── dataType_date
│       ├── image
│       └── csv
├── interim      <- Intermediate data that has been transformed.
│   └── ${rawDir}_howToProcess
│       ├── image
│       └── csv
├── processed    <- The final, canonical data sets for modeling.
│   └── ${rawDir}_howToProcess
│       ├── image
│       └── csv
└── external     <- Data from third party sources.
    └── dataType_date
        ├── image
        └── csv

# ls(s3)
$ aws s3 ls s3://myBucket/data/raw/

# download(s3 -> local)
$ aws s3 cp s3://myBucket/data/raw/ ./data/raw/ --recursive

# upload(local -> s3)
$ aws s3 cp ./data/processed/ s3://myBucket/data/processed/  --recursive

1. 前処理、後処理、そしてネットワークモデルの入出力を確認
2. Datasetの作成
3. DataLoaderの作成
4. ネットワークモデルの作成
5. 順伝播（forward）の定義
6. 損失関数の定義
7. 最適化手法の設定
8. 学習・検証の実施
9. テストデータで推論

# Datasetの作成
import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='../data/raw', train=True,
                                        download=True, transform=transform)

testset = torchvision.datasets.CIFAR10(root='../data/raw', train=False,
                                       download=True, transform=transform)

# DataLoaderの作成
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

# ネットワークモデルの作成, 順伝播（forward）の定義
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

# 損失関数の定義
criterion = nn.CrossEntropyLoss()

# 最適化手法の設定
import torch.optim as optim

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 学習・検証の実施
for epoch in range(2):  # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
print('Finished Training')

import sagemaker

sagemaker_session = sagemaker.Session()

bucket = sagemaker_session.default_bucket() # ex. 'sagemaker-us-east-2-xxxxxxxxxxxx'

role = sagemaker.get_execution_role() # ex. arn:aws:iam::xxxxxxxxxxxx:role/service-role/...

# We use the sagemaker.Session.upload_data function to upload our datasets to an S3 location. 
# The return value inputs identifies the location -- we will use this later when we start the training job.
inputs = sagemaker_session.upload_data(path='../data/raw', bucket=bucket, key_prefix='data/raw/cifar10')

# main()
import argparse
import sagemaker_containers

# class Net(nn.Module)
import torch.nn as nn
import torch.nn.functional as F

# def _train(args)
import torch
import torchvision
import torchvision.transforms as transforms

import torch.optim as optim

# def _save_model(model, model_dir), def model_fn(model_dir)
import os

## Network, Forward
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def _train(args):
    ## Dataset 
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    trainset = torchvision.datasets.CIFAR10(root=args.data_dir, train=True,
                                            download=True, transform=transform)

    testset = torchvision.datasets.CIFAR10(root=args.data_dir, train=False,
                                           download=True, transform=transform)

    ## DataLoader
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
                                              shuffle=True, num_workers=args.workers)

    testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size,
                                             shuffle=False, num_workers=args.workers)

    net = Net()

    ## Loss function
    criterion = nn.CrossEntropyLoss()

    ## Optimizer
    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum)

    ##Train and Validation
    for epoch in range(args.epochs):  # loop over the dataset multiple times
#     for epoch in range(1):  # loop over the dataset multiple times
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:    # print every 2000 mini-batches
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

    print('Finished Training')    
    return _save_model(net, args.model_dir)


def _save_model(model, model_dir):
    path = os.path.join(model_dir, 'model.pth')
    # recommended way from http://pytorch.org/docs/master/notes/serialization.html
    torch.save(model.cpu().state_dict(), path)

def model_fn(model_dir):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Net()
    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('--workers', type=int, default=2, metavar='W',
                        help='number of data loading workers (default: 2)')
    parser.add_argument('--epochs', type=int, default=2, metavar='E',
                        help='number of total epochs to run (default: 2)')
    parser.add_argument('--batch_size', type=int, default=4, metavar='BS',
                        help='batch size (default: 4)')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                        help='initial learning rate (default: 0.001)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='momentum (default: 0.9)')
    parser.add_argument('--dist_backend', type=str, default='gloo', help='distributed backend (default: gloo)')

    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)

    _train(parser.parse_args())

import subprocess

# instance_type = 'local'
instance_type = 'ml.m5.large'

if subprocess.call('nvidia-smi') == 0:
    ## Set type to GPU if one is present
    instance_type = 'local_gpu'

print("Instance type = " + instance_type)

from sagemaker.pytorch import PyTorch
from datetime import datetime, timedelta, timezone

# estimator
hyper_param = {
    'workers': 2,
    'epochs':2,
    'batch_size': 4,
    'lr': 0.001,
    'momentum': 0.9,
}

cifar10_estimator = PyTorch(entry_point='../models/cifar10_sagemaker_org.py',
                            hyperparameters=hyper_param,
                            role=role,
                            framework_version='1.1.0',
                            train_instance_count=1,
                            train_instance_type=instance_type)

# fit
# date = datetime.now(timezone(timedelta(hours=+9), 'JST')).strftime("%Y%m%d%H%M")

cifar10_estimator.fit(
    inputs=inputs,
#     job_name='' + date,
)

from sagemaker.pytorch import PyTorch
from datetime import datetime, timedelta, timezone
from sagemaker.tuner import IntegerParameter, ContinuousParameter, HyperparameterTuner

# estimator
hyper_param = {
    'workers': 2,
    'epochs':2,
    'batch_size': 4,
    'lr': 0.001,
    'momentum': 0.9,
}

cifar10_estimator = PyTorch(entry_point='../models/cifar10_sagemaker_org.py',
                            hyperparameters=hyper_param,
                            role=role,
                            framework_version='1.1.0',
                            train_instance_count=1,
                            train_instance_type=instance_type)

# tuner
hyperparameter_ranges = {
    'batch_size': IntegerParameter(4, 64),
    'lr': ContinuousParameter(1e-4, 0.1),
    'momentum': ContinuousParameter(0.5, 0.9)
}
metric_definitions = [
    {'Name': 'loss', 'Regex': 'loss: (\S+)'}
]

cifar10_tuner = HyperparameterTuner(estimator=cifar10_estimator,
                                    objective_metric_name='loss',
                                    objective_type='Minimize',
                                    hyperparameter_ranges=hyperparameter_ranges,
                                    metric_definitions=metric_definitions,
                                    max_jobs=1,
                                    max_parallel_jobs=1,
                                    early_stopping_type='Auto')

# fit
# date = datetime.now(timezone(timedelta(hours=+9), 'JST')).strftime("%Y%m%d%H%M")

cifar10_tuner.fit(
    inputs=inputs
#     job_name='' + date,
)

from sagemaker.pytorch import PyTorchModel

cifar10_predictor = cifar10_estimator.deploy(initial_instance_count=1,
                                             instance_type=instance_type)

cifar10_estimator.delete_endpoint()

# テストデータ準備
import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 画像表示準備
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 推論実行（サンプル）
dataiter = iter(testloader)
images, labels = dataiter.next()

imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%4s' % classes[labels[j]] for j in range(4)))

outputs = cifar10_predictor.predict(images.numpy())

_, predicted = torch.max(torch.from_numpy(np.array(outputs)), 1)

print('Predicted: ', ' '.join('%4s' % classes[predicted[j]]
                              for j in range(4)))

# 推論実行（精度算出）
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = cifar10_predictor.predict(images.numpy())
        _, predicted = torch.max(torch.from_numpy(np.array(outputs)), 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))