{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compile and Deploy the pretrained PyTorch model from model zoo with SageMaker Neo\n",
    "\n",
    "---\n",
    "\n",
    "***[주의] 본 모듈은 PyTorch EIA 1.3.1 버전이나 1.5.1 버전에서 훈련을 수행한 모델만 배포가 가능합니다. 코드가 정상적으로 수행되지 않는다면, 프레임워크 버전을 동일 버전으로 맞춰 주시기 바랍니다.***\n",
    "\n",
    "본 모듈에서는 Elastic Inference Accelerator(EIA)를 사용하여 모델을 배포해 보겠습니다.\n",
    "\n",
    "## Elastic Inference Accelerator\n",
    "훈련 인스턴스와 달리 실시간 추론 인스턴스는 계속 상시로 띄우는 경우가 많기에, 딥러닝 어플리케이션에서 low latency를 위해 GPU 인스턴스를 사용하면 많은 비용이 발생합니다.\n",
    "\n",
    "Amazon Elastic Inference는 저렴하고 메모리가 작은 GPU 기반 가속기를 Amazon EC2, Amazon ECS, Amazon SageMaker에 연결할 수 있는 서비스로, Accelerator가 CPU 인스턴스에 프로비저닝되고 연결됩니다. EIA를 사용하면 GPU 인스턴스에 근접한 퍼포먼스를 보이면서 인스턴스 실행 비용을 최대 75%까지 절감할 수 있습니다. \n",
    "\n",
    "모든 Amazon SageMaker 인스턴스 유형, EC2 인스턴스 유형 또는 Amazon ECS 작업을 지원하며, 대부분의 딥러닝 프레임워크를 지원하고 있습니다. 지원되는 프레임워크 버전은 AWS CLI로 확인할 수 있습니다.\n",
    "\n",
    "```bash\n",
    "aws ecr list-images --repository-name tensorflow-inference-eia --registry-id 763104351884\n",
    "aws ecr list-images --repository-name pytorch-inference-eia --registry-id 763104351884\n",
    "aws ecr list-images --repository-name mxnet-inference-eia --registry-id 763104351884\n",
    "```\n",
    "\n",
    "- 참조: https://aws.amazon.com/ko/blogs/korea/amazon-elastic-inference-gpu-powered-deep-learning-inference-acceleration/\n",
    "\n",
    "## TorchScript Compile (Tracing)\n",
    "\n",
    "PyTorch 프레임워크에서 EI를 사용하기 위해서는 [TorchScript](https://pytorch.org/docs/1.3.1/jit.html)로 모델을 컴파일해야 하며, 2022년 2월 시점에서는 PyTorch 1.3.1와 PyTorch 1.5.1을 지원하고 있습니다. TorchScript는 PyTorch 코드에서 직렬화 및 최적화 가능한 모델로 컴파일하며 Python 인터프리터의 글로벌 인터프리터 잠금 (GIL)과 무관하기 때문에 Python 외의 언어에서 로드 가능하고  최적화가 용이합니다.\n",
    "\n",
    "TorchScript로 변환하는 방법은 **tracing** 방식과 **scripting** 방식이 있으며, 본 핸즈온에서는 tracing 방식을 사용하겠습니다.\n",
    "참고로 tracing 방식은 샘플 입력 데이터를 모델에 입력 후 그 입력의 흐름(feedforward)을 기록하여 포착하는 메커니즘이며, scripting 방식은 모델 코드를 직접 분석해서 컴파일하는 방식입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, sagemaker, random\n",
    "\n",
    "import torch\n",
    "print(sagemaker.__version__)\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 1. Inference script\n",
    "---\n",
    "\n",
    "아래 코드 셀은 `src` 디렉토리에 SageMaker 추론 스크립트를 저장합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile src/infer_pytorch_eia.py\n",
    "\n",
    "import io\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image  # Training container doesn't have this package\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "    \n",
    "# To use new EIA inference API, customer should use attach_eia(model, eia_ordinal_number)\n",
    "VERSIONS_USE_NEW_API = [\"1.5.1\"]\n",
    "\n",
    "def model_fn(model_dir):\n",
    "    try:\n",
    "        loaded_model = torch.jit.load(\"model.pth\", map_location=torch.device(\"cpu\"))\n",
    "        if torch.__version__ in VERSIONS_USE_NEW_API:\n",
    "            import torcheia\n",
    "\n",
    "            loaded_model = loaded_model.eval()\n",
    "            loaded_model = torcheia.jit.attach_eia(loaded_model, 0)\n",
    "        return loaded_model\n",
    "    except Exception as e:\n",
    "        logger.exception(f\"Exception in model fn {e}\")\n",
    "        return None\n",
    "\n",
    "def transform_fn(model, payload, request_content_type='application/octet-stream', \n",
    "                 response_content_type='application/json'):\n",
    "\n",
    "    logger.info('Invoking user-defined transform function')\n",
    "\n",
    "    if request_content_type != 'application/octet-stream':\n",
    "        raise RuntimeError(\n",
    "            'Content type must be application/octet-stream. Provided: {0}'.format(request_content_type))\n",
    "\n",
    "    # preprocess\n",
    "    decoded = Image.open(io.BytesIO(payload))\n",
    "    preprocess = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(\n",
    "            mean=[\n",
    "                0.485, 0.456, 0.406], std=[\n",
    "                0.229, 0.224, 0.225]),\n",
    "    ])\n",
    "    normalized = preprocess(decoded)\n",
    "    batchified = normalized.unsqueeze(0)\n",
    "\n",
    "    # predict\n",
    "    # With EI, client instance should be CPU for cost-efficiency. Subgraphs with unsupported arguments run locally. Server runs with CUDA\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    batchified = batchified.to(device)\n",
    "    \n",
    "    # Please make sure model is loaded to cpu and has been eval(), in this example, we have done this step in model_fn()\n",
    "    with torch.no_grad():\n",
    "        if torch.__version__ in VERSIONS_USE_NEW_API:\n",
    "            # Please make sure torcheia has been imported\n",
    "            import torcheia\n",
    "\n",
    "            # We need to set the profiling executor for EIA\n",
    "            torch._C._jit_set_profiling_executor(False)\n",
    "            with torch.jit.optimized_execution(True):\n",
    "                result =  model.forward(batchified)\n",
    "        # Set the target device to the accelerator ordinal\n",
    "        else:\n",
    "            with torch.jit.optimized_execution(True, {\"target_device\": \"eia:0\"}):\n",
    "                result = model(batchified)\n",
    "\n",
    "    # Softmax (assumes batch size 1)\n",
    "    result = np.squeeze(result.detach().cpu().numpy())\n",
    "    result_exp = np.exp(result - np.max(result))\n",
    "    result = result_exp / np.sum(result_exp)\n",
    "\n",
    "    response_body = json.dumps(result.tolist())\n",
    "\n",
    "    return response_body, response_content_type"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 2. 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",
    "\n",
    "input_shape = [1,3,224,224]\n",
    "traced_model = torch.jit.trace(model.float().eval(), torch.zeros(input_shape).float())\n",
    "torch.jit.save(traced_model, 'model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local Inference without Endpoint\n",
    "\n",
    "충분한 검증 및 테스트 없이 훈련된 모델을 곧바로 실제 운영 환경에 배포하기에는 많은 위험 요소들이 있기 때문에, 로컬 환경 상에서 추론을 수행하면서 디버깅하는 것을 권장합니다. 아래 코드 셀의 코드를 예시로 참조해 주세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_inference(img_path, predictor, show_img=True):\n",
    "    with open(img_path, mode='rb') as file:\n",
    "        payload = bytearray(file.read())\n",
    "\n",
    "    response = predictor.predict(payload)\n",
    "    result = json.loads(response.decode())\n",
    "    pred_cls_idx, pred_cls_str, prob = parse_result(result, show_img)\n",
    "    \n",
    "    return pred_cls_idx, pred_cls_str, prob \n",
    "\n",
    "\n",
    "def parse_result(result, img_path, show_img=True):\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 matplotlib.pyplot as plt\n",
    "        img = Image.open(img_path)\n",
    "        plt.figure()\n",
    "        fig, ax = plt.subplots(1, figsize=(10,10))\n",
    "        ax.imshow(img)\n",
    "        overlay_text = f'{pred_cls_str} {prob:.2f}%'\n",
    "        ax.text(20, 40, overlay_text, style='italic',\n",
    "                bbox={'facecolor': 'yellow', 'alpha': 0.5, 'pad': 10}, fontsize=20)\n",
    "\n",
    "    return pred_cls_idx, pred_cls_str, prob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "모델 배포가 완료되었으면, 저자가 직접 준비한 샘플 이미지들로 추론을 수행해 보겠습니다."
   ]
  },
  {
   "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.infer_pytorch_eia import transform_fn\n",
    "\n",
    "path = \"./samples\"\n",
    "img_list = os.listdir(path)\n",
    "img_path_list = [os.path.join(path, img) for img in img_list]\n",
    "\n",
    "#test_idx = random.randint(0, len(img_list)-1)\n",
    "test_idx = 0\n",
    "img_path = img_path_list[test_idx]\n",
    "\n",
    "with open(img_path, mode='rb') as file:\n",
    "    payload = bytearray(file.read())"
   ]
  },
  {
   "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 = 'metadata/imagenet1000_clsidx_to_labels.txt'\n",
    "label_map = get_label_map_imagenet(label_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "#model = torch.jit.load('model.pth')\n",
    "#model = model.to(device)\n",
    "\n",
    "response_body, _ = transform_fn(traced_model, payload)\n",
    "result = json.loads(response_body)\n",
    "parse_result(result, img_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 3. Compile the Model\n",
    "---\n",
    "\n",
    "### 모델 압축"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tarfile.open('model.tar.gz', 'w:gz') as f:\n",
    "    f.add('model.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import time\n",
    "import sagemaker\n",
    "from sagemaker.pytorch import PyTorchModel\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "sm_sess = sagemaker.Session()\n",
    "region = sm_sess.boto_region_name\n",
    "bucket = sm_sess.default_bucket()\n",
    "\n",
    "instance_type = \"ml.m5.large\"\n",
    "accelerator_type = \"ml.eia2.xlarge\"\n",
    "\n",
    "# TorchScript model\n",
    "tar_filename = \"model.tar.gz\"\n",
    "\n",
    "# You can also upload model artifacts to S3\n",
    "# print('Upload tarball to S3')\n",
    "# model_data = sagemaker_session.upload_data(path=tar_filename, bucket=bucket, key_prefix=prefix)\n",
    "model_data = tar_filename\n",
    "\n",
    "endpoint_name = (\n",
    "    \"mnist-ei-traced-{}-{}\".format(instance_type, accelerator_type)\n",
    "    .replace(\".\", \"\")\n",
    "    .replace(\"_\", \"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_model = PyTorchModel(\n",
    "    model_data=model_data,\n",
    "    role=role,\n",
    "    entry_point=\"infer_pytorch_eia.py\",\n",
    "    source_dir=\"src\",\n",
    "    framework_version=\"1.5.1\",\n",
    "    py_version=\"py3\",\n",
    "    sagemaker_session=sm_sess,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deployment\n",
    "\n",
    "instance_type과 accelerator_type만 EIA에 적절하게 변경해 주시면 되며, 로컬 모드 배포도 가능합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attach EI remotely\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "from sagemaker.serializers import IdentitySerializer\n",
    "\n",
    "# Function will exit before endpoint is finished creating\n",
    "predictor = pytorch_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    accelerator_type=accelerator_type,\n",
    "    endpoint_name=endpoint_name,\n",
    "    serializer=IdentitySerializer(content_type='application/octet-stream'),\n",
    "    deserializer=JSONDeserializer(),\n",
    "    wait=False\n",
    ")\n",
    "\n",
    "# # Attach EI locally\n",
    "# # Deploys the model to a local endpoint\n",
    "# predictor = pytorch_model.deploy(\n",
    "#     initial_instance_count=1,\n",
    "#     instance_type='local',\n",
    "#     accelerator_type='local_sagemaker_notebook',\n",
    "#     serializer=IdentitySerializer(content_type='application/octet-stream'),\n",
    "#     deserializer=JSONDeserializer()        \n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wait for the endpoint jobs to complete\n",
    "\n",
    "엔드포인트가 생성될 때까지 기다립니다. 약 5-10분의 시간이 소요됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "def make_endpoint_link(region, endpoint_name, endpoint_task):\n",
    "    \n",
    "    endpoint_link = f'<b><a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region={region}#/endpoints/{endpoint_name}\">{endpoint_task} Review Endpoint</a></b>'   \n",
    "    return endpoint_link \n",
    "        \n",
    "endpoint_link = make_endpoint_link(region, predictor.endpoint_name, '[Deploy EIA model]')\n",
    "\n",
    "display(HTML(endpoint_link))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_sess.wait_for_endpoint(predictor.endpoint_name, poll=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 4. Inference\n",
    "---\n",
    "\n",
    "모델 배포가 완료되었으면, 추론을 수행합니다."
   ]
  },
  {
   "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",
    "\n",
    "img_list = os.listdir(path)\n",
    "img_path_list = [os.path.join(path, img) for img in img_list]\n",
    "print(img_path_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_idx = random.randint(0, len(img_list)-1)\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",
    "response = predictor.predict(payload)\n",
    "pred_cls_idx, pred_cls_str, prob = parse_result(response, img_path, show_img=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "마지막으로 latency를 측정합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "start_time = time.time()\n",
    "num_tests = 20\n",
    "for _ in range(num_tests):\n",
    "    response = predictor.predict(payload)\n",
    "end_time = (time.time()-start_time)\n",
    "print(f'EIA optimized inference time is {(end_time/num_tests)*1000:.4f} ms (avg)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Endpoint Clean-up\n",
    "\n",
    "SageMaker Endpoint로 인한 과금을 막기 위해, 본 핸즈온이 끝나면 반드시 Endpoint를 삭제해 주시기 바랍니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()\n",
    "pytorch_model.delete_model()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_amazonei_pytorch_latest_p37",
   "language": "python",
   "name": "conda_amazonei_pytorch_latest_p37"
  },
  "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}