{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 사전 훈련된 PyTorch 모델을 SageMaker Neo로 컴파일 후 Inf1 인스턴스에 배포하기\n",
    "\n",
    "---\n",
    "\n",
    "이 노트북에서는 사전 훈련된 `mnasnet` 모델을 SageMaker Neo로 컴파일 후, Inf1 인스턴스로 배포합니다. SageMaker Neo로 컴파일한 모델은 \n",
    "클라우드와 엣지 디바이스 어디에서나 실행할 수 있습니다. 또한, Inf1 인스턴스는 머신 러닝 추론 애플리케이션을 지원하기 위해 구축된 추론 전용 인스턴스이며, AWS에서 설계 및 구축한 고성능 머신 러닝 추론 칩인 AWS Inferentia 칩이 최대 16개 내장되어 있습니다.\n",
    "\n",
    "2021년 1월 기준으로 torchvision의 사전 훈련 모델을 지원하고 있으며 클라우드 인스턴스와 엣지 디바이스에서 PyTorch 1.4.0을 지원하고 있으며, AWS Inferentia에서 PyTorch 1.5.1을 지원하고 있습니다.\n",
    "\n",
    "SageMaker Neo가 지원하는 인스턴스 타입, 하드웨어 및 딥러닝 프레임워크는 아래 링크를 참조해 주세요.\n",
    "- 클라우드 인스턴스: https://docs.aws.amazon.com/sagemaker/latest/dg/neo-supported-cloud.html\n",
    "- 엣지 디바이스: https://docs.aws.amazon.com/sagemaker/latest/dg/neo-supported-devices-edge.html\n",
    "\n",
    "또한, SageMaker 인프라 상에서 PyTorch로 추론을 수행하는 튜토리얼 노트북 예제는 [pytorch-serving-endpoint.ipynb](pytorch-serving-endpoint.ipynb) 를 참조해 주세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SageMaker SDK를 최신 버전으로 업그레이드합니다. 본 노트북은 SDK 2.x 버전 이상에서 구동해야 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, sagemaker\n",
    "!{sys.executable} -m pip install -qU \"sagemaker>=2.11.0\"\n",
    "print(sagemaker.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[주의]** 주피터 노트북에 설치된 PyTorch 및 TorchVision 버전은 각각 1.4.0과 0.5.0이어야 합니다. SageMaker 노트북 인스턴스 상에서 본 예제를 실행한다면 아래 코드 셀을 주석 처리하지 말고 그대로 실행해 주세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!~/anaconda3/envs/pytorch_p36/bin/pip install -qU torch==1.4.0 torchvision==0.5.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "# 1. 단일 코어로 모델 컴파일 후 배포\n",
    "---\n",
    "우선, 이 섹션에서는 컴파일의 기본 설정값인 단일 코어로 모델을 컴파일합니다.\n",
    "다음 섹션에서 컴파일러 옵션을 사용하여 다중 코어에 대한 모델을 컴파일하는 방법을 살펴 보겠습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference Script\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile src/inference_pytorch_inf1.py\n",
    "\n",
    "def input_fn(request_body, request_content_type):\n",
    "    import torch\n",
    "    import torchvision.transforms as transforms\n",
    "    from PIL import Image\n",
    "    import io\n",
    "    f = io.BytesIO(request_body)\n",
    "    input_image = Image.open(f).convert('RGB')\n",
    "    preprocess = transforms.Compose([\n",
    "      transforms.Resize(255),\n",
    "      transforms.CenterCrop(224),\n",
    "      transforms.ToTensor(),\n",
    "      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ])\n",
    "    input_tensor = preprocess(input_image)\n",
    "    input_batch = input_tensor.unsqueeze(0)\n",
    "    return input_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import pre-trained model from TorchVision\n",
    "\n",
    "본 예제는 TorchVision의 pre-trained 모델 중 MnasNet을 사용합니다.\n",
    "MnasNet은 정확도(accuracy)와 모바일 디바이스의 latency를 모두 고려한 강화학습 기반 NAS(neural architecture search)이며, TorchVision은 image classification에 최적화된 MNasNet-B1을 내장하고 있습니다. \n",
    "(참조 논문: https://arxiv.org/pdf/1807.11626.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "import tarfile\n",
    "\n",
    "model = models.mnasnet1_0(pretrained=True)\n",
    "input_shape = [1,3,224,224]\n",
    "trace = torch.jit.trace(model.float().eval(), torch.zeros(input_shape).float())\n",
    "trace.save('model.pth')\n",
    "\n",
    "with tarfile.open('model.tar.gz', 'w:gz') as f:\n",
    "    f.add('model.pth')\n",
    "    f.add('src/inference_pytorch_neo.py')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 모델 아티팩트 S3로 전송"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import time\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "sess = sagemaker.Session()\n",
    "region = sess.boto_region_name\n",
    "bucket = sess.default_bucket()\n",
    "\n",
    "compilation_job_name = name_from_base('TorchVision-MnasNet-Neo-Inf1')\n",
    "\n",
    "model_key = '{}/model/model.tar.gz'.format(compilation_job_name)\n",
    "model_path = 's3://{}/{}'.format(bucket, model_key)\n",
    "boto3.resource('s3').Bucket(bucket).upload_file('model.tar.gz', model_key)\n",
    "print(\"Uploaded model to s3:\")\n",
    "print(model_path)\n",
    "\n",
    "sm_client = boto3.client('sagemaker')\n",
    "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)\n",
    "print(\"Output path for compiled model:\")\n",
    "print(compiled_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 모델 컴파일\n",
    "모델을 Inf1 인스턴스에 배포하기 전, SageMaker Neo를 사용하여 모델을 컴파일하여 배포 대상 디바이스 및 인스턴스에 대한 성능을 최적화합니다.\n",
    "Inf1 인스턴스에 배포할 모델을 컴파일하기 위해 `target_instance_family='ml_inf1'`을 설정하세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "pytorch_model = PyTorchModel(model_data=model_path,\n",
    "                             role=role,\n",
    "                             source_dir='src',\n",
    "                             entry_point='inference_pytorch_inf1.py',\n",
    "                             framework_version='1.5.1',\n",
    "                             py_version='py3')\n",
    "\n",
    "neo_model = pytorch_model.compile(target_instance_family='ml_inf1',\n",
    "                                  input_shape={'input0':[1,3,224,224]},\n",
    "                                  output_path=compiled_model_path,\n",
    "                                  framework='pytorch',\n",
    "                                  framework_version='1.5.1',\n",
    "                                  role=role,\n",
    "                                  job_name=compilation_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 컴파일된 모델 배포\n",
    "\n",
    "컴파일 완료 후, 곧바로 SageMaker 엔드포인트에 배포합니다. SageMaker의 Inf1 인스턴스는 `ml.inf1.xlarge, ml.inf1.2xlarge, ml.inf1.6xlarge 및 ml.inf1.24xlarge`를 지원하고 있으며, 본 예제에서는 가장 저렴한 `ml.inf1.xlarge`를 사용합니다.\n",
    "\n",
    "SageMaker가 관리하는 배포 클러스터를 프로비저닝하고 추론 컨테이너를 배포하기 때문에, 추론 서비스를 시작하는 데에는 약 5~10분 정도 소요됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "predictor = neo_model.deploy(instance_type='ml.inf1.xlarge', initial_instance_count=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "# 2. Inference\n",
    "---\n",
    "\n",
    "모델 배포가 완료되었으면, 추론을 수행해 보겠습니다. COCO dataset 2017 Test 이미지를 몇 장 준비했습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "클래스 인덱스에 대응하는 클래스명을 매핑하기 위한 딕셔너리를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.utils import get_label_map_imagenet\n",
    "label_file = 'files/imagenet1000_clsidx_to_labels.txt'\n",
    "label_map = get_label_map_imagenet(label_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "from src.inference_pytorch_neo import transform_fn\n",
    "\n",
    "path = \"./images/imgcls_test\"\n",
    "img_list = os.listdir(path)\n",
    "img_path_list = [os.path.join(path, img) for img in img_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_result(result, show_img=True):\n",
    "    \n",
    "    result_exp = np.exp(result - np.max(result))\n",
    "    result = result_exp / np.sum(result_exp)\n",
    "\n",
    "    pred_cls_idx = np.argmax(result)\n",
    "    pred_cls_str = label_map[str(pred_cls_idx)]\n",
    "    prob = np.amax(result)*100\n",
    "    \n",
    "    if show_img:\n",
    "        import cv2\n",
    "        import matplotlib.pyplot as plt\n",
    "        im = cv2.imread(img_path, cv2.COLOR_BGR2RGB)\n",
    "        font = cv2.FONT_HERSHEY_COMPLEX_SMALL\n",
    "        cv2.putText(im, f'{pred_cls_str} {prob:.2f}%', (10,40), font, 1, (0, 255, 0), 2, cv2.LINE_AA)\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(im[:,:,::-1])    \n",
    "\n",
    "    return pred_cls_idx, pred_cls_str, prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "\n",
    "test_idx = 2\n",
    "img_path = img_path_list[test_idx]\n",
    "\n",
    "with open(img_path, mode='rb') as file:\n",
    "    payload = bytearray(file.read())\n",
    "    \n",
    "sm_runtime = boto3.Session().client('sagemaker-runtime')\n",
    "\n",
    "response = sm_runtime.invoke_endpoint(EndpointName=predictor.endpoint_name,\n",
    "                                      ContentType='application/x-image',\n",
    "                                      Body=payload)\n",
    "result = json.loads(response['Body'].read().decode())\n",
    "pred_cls_idx, pred_cls_str, prob = parse_result(result)\n",
    "\n",
    "print(response)\n",
    "print(f'Most likely class: {pred_cls_str}, {prob:.2f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "for _ in range(100):\n",
    "    rresponse = sm_runtime.invoke_endpoint(EndpointName=predictor.endpoint_name,\n",
    "                                      ContentType='application/x-image',\n",
    "                                      Body=payload)\n",
    "inf1_inference_time = (time.time()-start)\n",
    "print(f'inf1 optimized inference time is {inf1_inference_time:.4f} ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Endpoint Clean-up\n",
    "\n",
    "SageMaker Endpoint로 인한 과금을 막기 위해, Endpoint를 삭제합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.delete_endpoint(predictor.endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "# 3. 컴파일러 옵션을 사용하여 다중 코어용 모델 컴파일 후 배포\n",
    "---\n",
    "\n",
    "이 섹션에서는 `inf1.xlarge`에서 지원하는 4개의 코어를 모두 활용해 보겠습니다. 본 예제에서는 모델 정의 시\n",
    "`env={'NEURONCORE_GROUP_SIZES': '2', 'SAGEMAKER_MODEL_SERVER_WORKERS': '2'}` 로 지정하고 모델 컴파일 시 뉴런 코어 개수를 2로 조절합니다. (`compiler_options={\"num-neuroncores\": 2}`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import time\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "sess = sagemaker.Session()\n",
    "region = sess.boto_region_name\n",
    "bucket = sess.default_bucket()\n",
    "\n",
    "compilation_job_name = name_from_base('TorchVision-ResNet18-Neo-Inf1')\n",
    "\n",
    "model_key = '{}/model/model.tar.gz'.format(compilation_job_name)\n",
    "model_path = 's3://{}/{}'.format(bucket, model_key)\n",
    "boto3.resource('s3').Bucket(bucket).upload_file('model.tar.gz', model_key)\n",
    "print(\"Uploaded model to s3:\")\n",
    "print(model_path)\n",
    "\n",
    "sm_client = boto3.client('sagemaker')\n",
    "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)\n",
    "print(\"Output path for compiled model:\")\n",
    "print(compiled_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 모델 컴파일\n",
    "컴파일러 옵션에서 뉴런 코어 수를 2로 전달합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "pytorch_model = PyTorchModel(model_data=model_path,\n",
    "                             role=role,\n",
    "                             entry_point='resnet18.py',\n",
    "                             framework_version='1.5.1',\n",
    "                             py_version='py3',\n",
    "                             env={'NEURONCORE_GROUP_SIZES': '2', 'SAGEMAKER_MODEL_SERVER_WORKERS': '2'}\n",
    "                             )\n",
    "neo_model = pytorch_model.compile(target_instance_family='ml_inf1',\n",
    "                                  input_shape={'input0':[1,3,224,224]},\n",
    "                                  output_path=compiled_model_path,\n",
    "                                  framework='pytorch',\n",
    "                                  framework_version='1.5.1',\n",
    "                                  role=role,\n",
    "                                  job_name=compilation_job_name,\n",
    "                                  compiler_options={\"num-neuroncores\": 2})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 컴파일된 모델 배포\n",
    "\n",
    "SageMaker가 관리하는 배포 클러스터를 프로비저닝하고 추론 컨테이너를 배포하기 때문에, 추론 서비스를 시작하는 데에는 약 5~10분 정도 소요됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "predictor = neo_model.deploy(instance_type='ml.inf1.xlarge', initial_instance_count=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "\n",
    "test_idx = 2\n",
    "img_path = img_path_list[test_idx]\n",
    "\n",
    "with open(img_path, mode='rb') as file:\n",
    "    payload = bytearray(file.read())\n",
    "    \n",
    "sm_runtime = boto3.Session().client('sagemaker-runtime')\n",
    "\n",
    "response = sm_runtime.invoke_endpoint(EndpointName=predictor.endpoint_name,\n",
    "                                      ContentType='application/x-image',\n",
    "                                      Body=payload)\n",
    "result = json.loads(response['Body'].read().decode())\n",
    "pred_cls_idx, pred_cls_str, prob = parse_result(result)\n",
    "\n",
    "print(response)\n",
    "print(f'Most likely class: {pred_cls_str}, {prob:.2f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "for _ in range(100):\n",
    "    rresponse = sm_runtime.invoke_endpoint(EndpointName=predictor.endpoint_name,\n",
    "                                      ContentType='application/x-image',\n",
    "                                      Body=payload)\n",
    "inf1_inference_time = (time.time()-start)\n",
    "print(f'inf1 optimized inference time is {inf1_inference_time:.4f} ms')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.delete_endpoint(predictor.endpoint_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}