CR

Use new inference containers for both uncompiled and compiled flows.

Author: kkoppolu1

sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb

@@ -158,16 +158,6 @@
     "                                        )"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# TODO(kkoppolu): Delete after new SDK version sets the image URI correctly\n",
-    "compiled_model.image_uri = compiled_model.image_uri.replace(\"neo\", \"inference\")"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -251,13 +241,6 @@
    "source": [
     "sess.delete_endpoint(predictor.endpoint_name)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/code/vgg19_bn.py renamed to sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/code/vgg19_bn_compiled.py

@@ -46,7 +46,7 @@ def transform_fn(model, payload, request_content_type,
     result = model.forward(batchified)
 
     # Softmax (assumes batch size 1)
-    result = np.squeeze(result.cpu().numpy())
+    result = np.squeeze(result.cpu().detach().numpy())
     result_exp = np.exp(result - np.max(result))
     result = result_exp / np.sum(result_exp)
 

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/code/vgg19_bn_uncompiled.py

@@ -0,0 +1,73 @@
+import io
+import json
+import logging
+import os
+import pickle
+
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from PIL import Image  # Training container doesn't have this package
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+# ------------------------------------------------------------ #
+# Neo host methods                                             #
+# ------------------------------------------------------------ #
+
+def transform_fn(model, payload, request_content_type,
+                 response_content_type):
+
+    logger.info('Invoking user-defined transform function')
+
+    if request_content_type != 'application/octet-stream':
+        raise RuntimeError(
+            'Content type must be application/octet-stream. Provided: {0}'.format(request_content_type))
+
+    # preprocess
+    decoded = Image.open(io.BytesIO(payload))
+    preprocess = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[
+                0.485, 0.456, 0.406], std=[
+                0.229, 0.224, 0.225]),
+    ])
+    normalized = preprocess(decoded)
+    batchified = normalized.unsqueeze(0)
+    
+    # predict
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    batchified = batchified.to(device)
+    result = model.forward(batchified)
+
+    # Softmax (assumes batch size 1)
+    result = np.squeeze(result.cpu().detach().numpy())
+    result_exp = np.exp(result - np.max(result))
+    result = result_exp / np.sum(result_exp)
+
+    response_body = json.dumps(result.tolist())
+    content_type = 'application/json'
+
+    return response_body, content_type
+
+
+# ------------------------------------------------------------ #
+# Sagemaker Hosting methods                                    #
+# ------------------------------------------------------------ #
+
+
+def model_fn(model_dir):
+
+    logger.info('model_fn')
+    with torch.neo.config(model_dir=model_dir, neo_runtime=True):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # The compiled model is saved as "compiled.pt"
+        model = torch.jit.load(os.path.join(model_dir, 'model.pth'))
+        model = model.to(device)
+
+        return model

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb

@@ -93,6 +93,7 @@
     "import sagemaker\n",
     "import time\n",
     "from sagemaker.utils import name_from_base\n",
+    "from sagemaker import image_uris\n",
     "\n",
     "role = sagemaker.get_execution_role()\n",
     "sess = sagemaker.Session()\n",
@@ -106,9 +107,11 @@
     "\n",
     "data_shape = '{\"input0\":[1,3,224,224]}'\n",
     "target_device = 'ml_c5'\n",
-    "framework = 'PYTORCH'\n",
-    "framework_version = '1.2.0'\n",
-    "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)"
+    "framework = 'pytorch'\n",
+    "framework_version = '1.4.0'\n",
+    "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)\n",
+    "\n",
+    "inference_image_uri = image_uris.retrieve(f'neo-{framework}', region, framework_version, instance_type=target_device)"
    ]
   },
   {
@@ -125,13 +128,17 @@
    "outputs": [],
    "source": [
     "from sagemaker.pytorch.model import PyTorchModel\n",
+    "from sagemaker.predictor import Predictor\n",
     "\n",
     "pt_vgg = PyTorchModel(model_data=model_path,\n",
     "                      framework_version=framework_version,\n",
-    "                      role=role,                               \n",
-    "                      entry_point='vgg19_bn_old.py',\n",
+    "                      predictor_cls=Predictor,\n",
+    "                      role=role,    \n",
     "                      sagemaker_session=sess,\n",
-    "                      py_version='py3'\n",
+    "                      entry_point='vgg19_bn_uncompiled.py',\n",
+    "                      source_dir='code',\n",
+    "                      py_version='py3',\n",
+    "                      image_uri=inference_image_uri\n",
     "                     )"
    ]
   },
@@ -176,7 +183,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "#### Image Pre-processing"
+    "#### Read the image payload"
    ]
   },
   {
@@ -185,19 +192,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
-    "from PIL import Image\n",
-    "from torchvision import transforms\n",
-    "import numpy as np\n",
-    "input_image = Image.open('cat.jpg')\n",
-    "preprocess = transforms.Compose([\n",
-    "    transforms.Resize(256),\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) # create a mini-batch as expected by the model"
+    "import json\n",
+    "\n",
+    "with open('cat.jpg', 'rb') as f:\n",
+    "    payload = f.read()\n",
+    "    payload = bytearray(payload) "
    ]
   },
   {
@@ -216,7 +215,7 @@
     "import time\n",
     "start = time.time()\n",
     "for _ in range(1000):\n",
-    "    output = vgg_predictor.predict(input_batch)\n",
+    "    output = vgg_predictor.predict(payload)\n",
     "inference_time = (time.time()-start)\n",
     "print('Inference time is ' + str(inference_time) + 'millisecond')"
    ]
@@ -227,7 +226,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "_, predicted = torch.max(torch.from_numpy(np.array(output)), 1)"
+    "import numpy as np\n",
+    "result = json.loads(output.decode())\n",
+    "predicted = np.argmax(result)"
    ]
   },
   {
@@ -250,7 +251,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"Result: label - \" + object_categories[str(predicted.item())])"
+    "print(\"Result: label - \" + object_categories[str(predicted)])"
    ]
   },
   {
@@ -277,39 +278,6 @@
     "## Neo optimization"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Update framework version"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "framework_version = '1.4.0'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Re-create the model archive"
-   ]
-  },
-  {
-   "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": "markdown",
    "metadata": {},
@@ -331,7 +299,7 @@
     "                               framework_version = framework_version,\n",
     "                               role=role,\n",
     "                               sagemaker_session=sess,\n",
-    "                               entry_point='vgg19_bn.py',\n",
+    "                               entry_point='vgg19_bn_compiled.py',\n",
     "                               source_dir='code',\n",
     "                               py_version='py3',\n",
     "                               env={'MMS_DEFAULT_RESPONSE_TIMEOUT': '500'}\n",
@@ -361,16 +329,6 @@
     "                                        )"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# TODO(kkoppolu): Delete after new SDK version sets the image URI correctly\n",
-    "compiled_model.image_uri = compiled_model.image_uri.replace(\"neo\", \"inference\")"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -382,19 +340,6 @@
     "                                 )"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json\n",
-    "\n",
-    "with open('cat.jpg', 'rb') as f:\n",
-    "    payload = f.read()\n",
-    "    payload = bytearray(payload) "
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -435,13 +380,6 @@
    "source": [
     "sess.delete_endpoint(predictor.endpoint_name)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {