{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a224550f",
   "metadata": {},
   "source": [
    "# Amazon SageMaker 持ち込みPyTorchコンテナによる MNIST の学習（Studioノートブック用）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc0ca35",
   "metadata": {},
   "source": [
    "### 1.背景\n",
    "\n",
    "MNISTは、手書き文字の分類に広く使用されているデータセットです。 70,000個のラベル付きの28x28ピクセルの手書き数字のグレースケール画像で構成されています。 データセットは、60,000個のトレーニング画像と10,000個のテスト画像に分割されます。 手書きの数字 0から9の合計10のクラスがあります。 \n",
    "\n",
    "このチュートリアルでは、SageMaker のマネージドコンテナイメージ以外に、独自のコンテナイメージを持ち込める持ち込みのことを学習していただくために、持ち込みの独自コンテナを使用して SageMaker で MNIST モデルをトレーニングする方法を示します。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26343f3",
   "metadata": {},
   "source": [
    "### 2.セットアップ\n",
    "\n",
    "SageMaker セッションを作成し、設定を開始し\n",
    "Dockerfileから、持ち込み用のコンテンナを作成していきます。\n",
    "Dockerfileはcontainerディレクトリ配下に配置されていて、PyTorchコンテナにrequirements.txt(pip installするパッケージリスト)経由でSageMaker Training Toolkitをpipインストールするように記述されています。\n",
    "\n",
    "SageMaker Training Toolkitは、どのDockerコンテナにも簡単に追加でき、SageMakerと互換性のあるモデルをトレーニングすることができるパッケージです。詳細は下記をご確認ください。\n",
    "https://github.com/aws/sagemaker-training-toolkit\n",
    "\n",
    "最後にdocker buildコマンドによりビルドしていきます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d06e6afb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68781009",
   "metadata": {},
   "source": [
    "Amazon SageMaker Studio ではDockerコマンドが環境上使えなくなっていますので\n",
    "コンテナイメージのビルドの際はdockerコマンドの代わりにsm-docker build コマンドを使う必要があります。\n",
    "sm-docker build を実行する際は、sm-docker build が裏側で AWS CodeBuild を使うようになっていてるため\n",
    "StudioノートにアタッチしたIAMロールに以下の作業を追加で実施いただく必要があります\n",
    "\n",
    " 1.”AmazonEC2ContainerRegistryPowerUser”の追加アタッチ\n",
    "\n",
    " 2.信頼ポリシーの追加（後述）\n",
    "\n",
    " 3.インラインポリシーの追加（後述）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c0195ce",
   "metadata": {},
   "source": [
    "sm-docker のインストール"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970f741d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker-studio-image-build"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44141393",
   "metadata": {},
   "source": [
    "Studioノートブックにアタッチされているロールに以下のコマンドの出力結果をコピーして追加の信頼関係を追加してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e4bf02",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat container/trust_relationships.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88eb3472",
   "metadata": {},
   "source": [
    "Studioノートブックにアタッチされているロールに以下のコマンドの出力結果をコピーしてインラインポリシー（もしくはポリシーを別途作成して）アタッチしてください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd648298",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat container/inline_policy.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40c91139",
   "metadata": {},
   "source": [
    "### 3.Amazon ECRリポジトリの作成とDockerImageの登録\n",
    "\n",
    "sm-docker buildコマンドでAmazon ECRにsagemaker-toolkit-container-pytorchというレポジトリを作成し、Studioノートブックに作成したDockerImageをpushしていきます。\n",
    "\n",
    "sm-docker buildコマンドを実行するとrepositoryオプションに指定したレポジトリとタグ名でAmazon ECRにイメージがpushされます。\n",
    "\n",
    "詳細は下記をご確認ください\n",
    "https://pypi.org/project/sagemaker-studio-image-build/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71509c99",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "IMAGE_NAME = 'toolkit-container-pytorch'\n",
    "TAG=':1.11'\n",
    "REGION = boto3.session.Session().region_name\n",
    "\n",
    "MY_ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account')\n",
    "MY_ECR_ENDPOINT = f'{MY_ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/'\n",
    "\n",
    "MY_REPOSITORY_NAME = f'sagemaker-{IMAGE_NAME}'\n",
    "MY_REPOSITORY_URI = f'{MY_ECR_ENDPOINT}sagemaker-{IMAGE_NAME}'\n",
    "MY_IMAGE_URI = f'{MY_REPOSITORY_URI}{TAG}'\n",
    "\n",
    "%cd ./container\n",
    "!sm-docker build . --repository {MY_REPOSITORY_NAME}{TAG}\n",
    "%cd ../\n",
    "\n",
    "print(f'コンテナイメージは {MY_IMAGE_URI} へ登録されています。')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4211fa2f",
   "metadata": {},
   "source": [
    "### 4.セットアップ\n",
    "\n",
    "学習およびモデルデータに使用する S3 バケットとプレフィックスは、ノートブックインスタンス、トレーニング、およびホスティングと同じリージョン内にある必要があります。\n",
    "\n",
    "データへの学習およびホスティングアクセスを提供するために使用される IAM ロール arn を用います。 ノートブックインスタンス、学習インスタンス、および/またはホスティングインスタンスに複数のロールが必要な場合は、 sagemaker.get_execution_role（） を、適切な IAM ロール arn 文字列に置き換えてください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ecf4633",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker.pytorch import PyTorch\n",
    "bucket = sagemaker.session.Session().default_bucket()\n",
    "prefix = 'sagemaker/DEMO-pytorch-mnist-byoc'\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "099d3164",
   "metadata": {},
   "source": [
    "### 5.データの取得"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66674b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp s3://fast-ai-imageclas/mnist_png.tgz . --no-sign-request\n",
    "!tar -xzf mnist_png.tgz -C /opt/ml --no-same-owner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f02fa304",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "import os\n",
    "\n",
    "root_dir_studio = '/opt/ml'\n",
    "data_dir = os.path.join(root_dir_studio,'data') \n",
    "training_dir = os.path.join(root_dir_studio,'mnist_png/training') \n",
    "test_dir = os.path.join(root_dir_studio,'mnist_png/testing') \n",
    "\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "training_data = datasets.ImageFolder(root=training_dir,\n",
    "                            transform=transforms.Compose([\n",
    "                            transforms.Grayscale(),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize((0.1307,), (0.3081,))]))\n",
    "test_data = datasets.ImageFolder(root=test_dir,\n",
    "                            transform=transforms.Compose([\n",
    "                            transforms.Grayscale(),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize((0.1307,), (0.3081,))]))\n",
    "\n",
    "training_data_loader = DataLoader(training_data, batch_size=len(training_data))\n",
    "training_data_loaded = next(iter(training_data_loader))\n",
    "torch.save(training_data_loaded, os.path.join(data_dir, 'training.pt'))\n",
    "\n",
    "test_data_loader = DataLoader(test_data, batch_size=len(test_data))\n",
    "test_data_loaded = next(iter(test_data_loader))\n",
    "torch.save(test_data_loaded, os.path.join(data_dir, 'test.pt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "462f4864",
   "metadata": {},
   "source": [
    "### 6.データをS3にアップロードする\n",
    "\n",
    "データセットを S3 にアップロードするには、 sagemaker.Session.upload_data 関数を使用します。 戻り値として入力した S3 のロケーションは、後で学習ジョブを実行するときに使用します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11383e53",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = sagemaker.session.Session().upload_data(path=data_dir, bucket=bucket, key_prefix=prefix)\n",
    "print('input spec (in this case, just an S3 path): {}'.format(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a6d7648",
   "metadata": {},
   "source": [
    "### 7.学習を開始する\n",
    "\n",
    "学習の条件を設定するため、Estimator クラスの子クラスの PyTorch オブジェクトを作成します。 ここでは、PyTorchスクリプト、IAMロール、および（ジョブごとの）ハードウェア構成を渡す PyTorch Estimator を定義しています。また合わせてentry_point を指定することで、依存するスクリプト群をコンテナにコピーして、学習時に使用することが可能です。\n",
    "\n",
    "トレーニングジョブを開始する際、Amazon S3のどこにソースが格納されているか、起動するモジュールは何かをsagemaker-training-toolkitライブラリに知らせる必要があります。今回はPython SDKのフレームワーク使用していますので自動的に設定されています。 Estimator Classを利用する際は必要になりますのでご注意ください。\n",
    "\n",
    "最後に、Amazon SageMaker Python SDKで定義されている汎用Estimatorオブジェクトのfit()メソッドを呼び出すことで、学習ジョブを実行することができます"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9c2daaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = PyTorch(\n",
    "    entry_point='train.py',\n",
    "    image_uri=MY_IMAGE_URI,\n",
    "    role=sagemaker.get_execution_role(),\n",
    "    hyperparameters={\n",
    "        'lr':0.01,\n",
    "        'batch-size':16\n",
    "    },\n",
    "    instance_count=1,\n",
    "    instance_type='ml.g4dn.xlarge',\n",
    ")\n",
    "estimator.fit({'training': inputs})\n",
    "\n",
    "#Estimatorクラスで学習する場合\n",
    "#estimator = sagemaker.estimator.Estimator(\n",
    "#    image_uri=MY_IMAGE_URI,\n",
    "#    role=sagemaker.get_execution_role(),\n",
    "#    hyperparameters={\n",
    "#        'lr':0.01,\n",
    "#        'batch-size':16,\n",
    "#        'sagemaker_program' : 'train.py',\n",
    "#        'sagemaker_submit_directory' : 's3://'+bucket+'/'+estimator._current_job_name+'/source/sourcedir.tar.gz'\n",
    "#    },\n",
    "#    instance_count=1,\n",
    "#    instance_type='ml.g4dn.xlarge',\n",
    "#)\n",
    "#estimator.fit({'training': inputs})"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.m5.xlarge",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.8 Python 3.6 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:ap-northeast-1:102112518831:image/1.8.1-cpu-py36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}