{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Compile the pretrained PyTorch model with SageMaker Neo for On-device\n", "---\n", "\n", "이 노트북에서는 사전 훈련된 MnasNet 기반 이미지 분류(Image classification) 모델을 SageMaker Neo로 컴파일하여 배포합니다. SageMaker Neo는 머신 러닝 모델을 하드웨어에 맞게 최적화하는 API로, Neo로 컴파일한 모델은 클라우드와 엣지 디바이스 어디에서나 실행할 수 있습니다.\n", "\n", "SageMaker Neo에서 지원하는 인스턴스 유형, 하드웨어 및 딥러닝 프레임워크는 아래 링크를 참조하세요.\n", "(본 예제 코드는 2021년 2월 기준으로 작성되었으며, 작성 시점에서 PyTroch 1.8.0까지 지원하고 있습니다. 단, AWS Inferentia 기반 인스턴스로 배포 시에는 PyTorch 1.7.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" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%load_ext autoreload\n", "%autoreload 2\n", "%store -r\n", "%store" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import logging, sys\n", "def _get_logger():\n", " '''\n", " # https://stackoverflow.com/questions/17745914/python-logging-module-is-printing-lines-multiple-times\n", " '''\n", " loglevel = logging.DEBUG\n", " l = logging.getLogger(__name__)\n", " if not l.hasHandlers():\n", " l.setLevel(loglevel)\n", " logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) \n", " l.handler_set = True\n", " return l \n", "\n", "logger = _get_logger()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os, sys, sagemaker\n", "sys.path.insert(0, \"./src\")\n", "print(sagemaker.__version__)\n", "model_trace_name = 'model.pth'\n", "sample_img_path = \"samples\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "## 1. Inference script\n", "---\n", "\n", "아래 코드 셀은 SageMaker 추론 스크립트를 `src` 디렉토리에 저장합니다." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%writefile src/infer_pytorch_neo.py\n", "\n", "import io\n", "import json\n", "import logging\n", "import os\n", "import pickle\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", "\n", "def model_fn(model_dir):\n", " import neopytorch\n", "\n", " logger.info(\"model_fn\")\n", " neopytorch.config(model_dir=model_dir, neo_runtime=True)\n", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " # The compiled model is saved as \"compiled.pt\"\n", " model = torch.jit.load(os.path.join(model_dir, \"compiled.pt\"), map_location=device)\n", "\n", " # It is recommended to run warm-up inference during model load\n", " sample_input_path = os.path.join(model_dir, \"sample_input.pkl\")\n", " with open(sample_input_path, \"rb\") as input_file:\n", " model_input = pickle.load(input_file)\n", " if torch.is_tensor(model_input):\n", " model_input = model_input.to(device)\n", " model(model_input)\n", " elif isinstance(model_input, tuple):\n", " model_input = (inp.to(device) for inp in model_input if torch.is_tensor(inp))\n", " model(*model_input)\n", " else:\n", " print(\"Only supports a torch tensor or a tuple of torch tensors\")\n", "\n", " return model\n", " \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", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " batchified = batchified.to(device)\n", " result = model.forward(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": [ "
\n", "\n", "## 2. Load trained model\n", "---\n", "\n", "사전 훈련된 모델을 로드합니다. 다른 프레임워크 버전과의 호환성 문제 및 직렬화 중 문제를 줄이기 위해 가능한 한 전체 모델을 로드하는 것보다 모델 구조를 먼저 초기화하고 모델 가중치(weight)를 로드하는 것을 권장합니다." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torchvision.models as models\n", "import tarfile\n", "import src.infer_utils as infer_utils\n", "\n", "classes_dict = infer_utils.load_classes_dict('classes_dict_imagenet.json')\n", "num_classes = len(classes_dict)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Detect if we have a GPU available\n", "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n", "\n", "model = models.mnasnet1_0(pretrained=True)\n", "#model = models.mobilenet_v2(pretrained=True)\n", "model = model.to(device)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torchvision.models as models\n", "import tarfile\n", "\n", "input_shape = [1,3,224,224]\n", "dummy_input = torch.zeros(input_shape).float()\n", "dummy_input = dummy_input.to(device)\n", "trace = torch.jit.trace(model.float().eval(), dummy_input)\n", "trace.save(model_trace_name)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Local Inference without Endpoint\n", "\n", "충분한 검증 및 테스트 없이 훈련된 모델을 곧바로 실제 운영 환경에 배포하기에는 많은 위험 요소들이 있기 때문에, 로컬 환경 상에서 추론을 수행하면서 디버깅하는 것을 권장합니다. 아래 코드 셀의 코드를 예시로 참조해 주세요." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import random\n", "import numpy as np\n", "from io import BytesIO\n", "from PIL import Image\n", "import src.infer_utils as infer_utils\n", "from src.infer_pytorch_neo import transform_fn\n", "\n", "model = torch.jit.load(model_trace_name)\n", "model = model.to(device)\n", "\n", "img_list = os.listdir(sample_img_path)\n", "img_path_list = [os.path.join(sample_img_path, img) for img in img_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_body, _ = transform_fn(model, payload)\n", "result = json.loads(response_body)\n", "infer_utils.parse_result(result, classes_dict, img_path, show_img=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "## 3. Compile Model with SageMaker Neo\n", "---\n", "\n", "### Overview\n", "\n", "Neo-AI는 다양한 머신 러닝 프레임워크를 지원하며 정확도 손실을 최소화하면서 자동으로 모델을 최적화합니다. Neo-AI 컴파일러는 타겟 디바이스의 OS 및 하드웨어 플랫폼에 맞게 모델을 자동으로 최적화하고 딥러닝 런타임에서 모델을 실행 가능한 형태로 변환합니다. 딥러닝 런타임은 머신 러닝 프레임워크와 엣지 디바이스에 상관없이 단 두 줄의 코드로 추론을 수행할 수 있으며 런타임 버전은 지속적으로 업데이트됩니다.\n", "\n", "그리고 AWS 계정이 있다면 Neo-AI 기반의 관리형 서비스인 Amazon SageMaker Neo를 사용할 수 있습니다. SageMaker Neo는 간단한 API 호출이나 UI로 추가 패키지나 인프라 설정 및 요금 부과 없이 동시에 여러 타켓 디바이스들에 적합한 모델을 컴파일할 수 있습니다." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "with tarfile.open('model.tar.gz', 'w:gz') as f:\n", " f.add(model_trace_name)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Model Compilation\n", "\n", "아래 코드 셀은 배포 유형에 따른 다양한 유즈케이스를 고려하여, 아래의 5가지 유즈케이스에 대한 컴파일 job을 동시에 시작합니다.\n", "\n", "- Cloud (CPU, `ml_m5` instance)\n", "- Cloud (CPU, `ml_c5` instance)\n", "- Cloud (GPU, `ml_g4dn` instance)\n", "- NVIDIA Jetson nano (CPU)\n", "- NVIDIA Jetson nano (GPU)\n", "\n", "NVIDIA Jetpack에 따라 디바이스의 CUDA 버전 또는 TensorRT 버전이 호환되지 않을 수 있으며, GPU 모델을 로드하는 데 수십 초가 걸리므로 CPU 모델을 함께 컴파일하고 테스트하는 것이 좋은 전략입니다.\n", "\n", "컴파일은 약 4-6분이 소요됩니다.\n", "\n", "**[Caution] 컴파일 중 오류가 발생하면 이 코드를 실행하는 노트북의 PyTorch 버전을 반드시 확인히세요. PyTorch 버전이 일치해야 합니다. 이 실습에서는 PyTorch 1.8을 사용합니다.**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import time, boto3, sagemaker\n", "role = sagemaker.get_execution_role()\n", "bucket = sagemaker.Session().default_bucket()\n", "\n", "# For cloud ML inference\n", "compilation_job_cloud_cpu_m5 = infer_utils.compile_model_for_cloud(\n", " role, bucket, target_device='ml_m5', dataset_dir=None, framework_version='1.8'\n", ")\n", "compilation_job_cloud_cpu_c5 = infer_utils.compile_model_for_cloud(\n", " role, bucket, target_device='ml_c5', dataset_dir=None, framework_version='1.8'\n", ")\n", "compilation_job_cloud_gpu = infer_utils.compile_model_for_cloud(\n", " role, bucket, target_device='ml_g4dn', dataset_dir=None, framework_version='1.8'\n", ")\n", "\n", "# For on-device ML inference\n", "compilation_job_jetson_cpu = infer_utils.compile_model_for_jetson(\n", " role, bucket, dataset_dir=None, use_gpu=False\n", ")\n", "compilation_job_jetson_gpu = infer_utils.compile_model_for_jetson(\n", " role, bucket, dataset_dir=None, use_gpu=True\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "compilation_jobs = [compilation_job_cloud_cpu_m5, compilation_job_cloud_cpu_c5, compilation_job_cloud_gpu, \n", " compilation_job_jetson_cpu, compilation_job_jetson_gpu]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "sm_client = boto3.client('sagemaker')\n", "\n", "max_time = time.time() + 15*60 # 15 mins\n", "for job in compilation_jobs:\n", " while time.time() < max_time:\n", " resp = sm_client.describe_compilation_job(CompilationJobName=job['job_name']) \n", " if resp['CompilationJobStatus'] in ['STARTING', 'INPROGRESS']:\n", " print('Running...')\n", " else:\n", " print(resp['CompilationJobStatus'], job)\n", " break\n", " time.sleep(30)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Review Compilation Jobs on AWS Console" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from IPython.core.display import display, HTML\n", "region = boto3.Session().region_name\n", "\n", "for job in compilation_jobs:\n", " job_name = job['job_name']\n", " display(\n", " HTML(\n", " 'Review Compilation Job for {}'.format(\n", " region, job_name, job_name\n", " )\n", " )\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Copy Compiled model to local\n", "\n", "아래 코드 셀은 컴파일된 모델을 S3에서 로컬로 복사합니다. 클라우드의 경우 인스턴스의 엔드포인트를 생성하여 실시간 배포가 가능하며, NVIDIA Jetson nano와 같은 온디바이스의 경우 모델을 디바이스에 복사하고 DLR을 설치합니다. DLR을 사용하면 PyTorch 및 TensorFlow와 같은 별도의 프레임워크를 설치할 필요 없이 간단한 API 호출로 모델을 쉽게 추론할 수 있습니다.\n", "\n", "- Installing DLR: https://neo-ai-dlr.readthedocs.io/en/latest/install.html" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model_root_path = 'neo-model'\n", "!rm -rf {model_root_path}\n", "for job in compilation_jobs:\n", " model_path = f\"{model_root_path}/{job['job_name']}\"\n", " os.makedirs(model_path, exist_ok=True)\n", " !aws s3 cp {job['s3_compiled_model_path']} {model_path} --recursive " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "## 4. On-device Deployment\n", "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "컴파일된 모델을 AWS IoT Greengrass 컴포넌트로 배포하는 경우 아래 코드 셀의 예시처럼 출력값을 기록해 두세요. 온디바이스에서 테스트하려면 아래 쉘 명령어의 예시처럼 실행하시면 됩니다.\n", "\n", "```shell\n", "rm -rf model_cpu\n", "mkdir model_cpu && cd model_cpu\n", "aws s3 cp [MODEL-CLOUD-CPU-S3-PATH] . --recursive\n", "tar -xzvf model-ml_m5.tar.gz && rm model-ml_m5.tar.gz\n", "```" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model_cloud_cpu_s3_path = compilation_jobs[0]['s3_compiled_model_path']\n", "print(model_cloud_cpu_s3_path)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "모델을 온디바이스로 복사한 후 압축을 해제하였다면, 아래 스크립트를 온디바이스로 복사&수정하여 추론을 수행해 봅니다." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "import logging, sys\n", "import cv2\n", "import glob\n", "import json\n", "import numpy as np\n", "import dlr\n", "from dlr import DLRModel\n", "\n", "\n", "def load_classes_dict(filename='classes_dict.json'):\n", " with open(filename, 'r') as fp:\n", " classes_dict = json.load(fp)\n", "\n", " classes_dict = {int(k):v for k,v in classes_dict.items()} \n", " return classes_dict\n", " \n", "\n", "def load_image(image_path):\n", " image_data = cv2.imread(image_path)\n", " image_data = cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)\n", " return image_data\n", "\n", "\n", "def preprocess_image(image, image_shape=(224,224)):\n", " cvimage = cv2.resize(image, image_shape)\n", " img = np.asarray(cvimage, dtype='float32')\n", " img /= 255.0 # scale 0 to 1\n", " mean = np.array([0.485, 0.456, 0.406]) \n", " std = np.array([0.229, 0.224, 0.225])\n", " img = (img - mean) / std\n", " img = np.transpose(img, (2,0,1)) \n", " img = np.expand_dims(img, axis=0) # e.g., [1x3x224x224]\n", " return img\n", "\n", "\n", "def softmax(x):\n", " x_exp = np.exp(x - np.max(x))\n", " f_x = x_exp / np.sum(x_exp)\n", " return f_x\n", "\n", "\n", "device = 'cpu'\n", "model = DLRModel(f'model_{device}', device)\n", "sample_image_dir = 'sample_images'\n", "classes_dict = load_classes_dict('classes_dict.json')\n", "\n", "extensions = (f\"{sample_image_dir}/*.jpg\", f\"{sample_image_dir}/*.jpeg\")\n", "img_filelist = [f for f_ in [glob.glob(e) for e in extensions] for f in f_]\n", "print(img_filelist)\n", "\n", "for img_filepath in img_filelist[:-1]:\n", " ground_truth = img_filepath.split('/')[-1]\n", " img = load_image(img_filepath)\n", " img_data = preprocess_image(img)\n", " \n", " output = model.run(img_data) \n", " probs = softmax(output[0][0])\n", " sort_classes_by_probs = np.argsort(probs)[::-1]\n", "\n", " idx = sort_classes_by_probs[0]\n", " print(\"+\"*80)\n", " print(f'predicted = {classes_dict[idx]}, {probs[idx]*100:.2f}%')\n", " print(f'ground_truth = {ground_truth}') \n", "``` " ] } ], "metadata": { "kernelspec": { "display_name": "conda_pytorch_p37", "language": "python", "name": "conda_pytorch_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.10" } }, "nbformat": 4, "nbformat_minor": 4 }