Merge branch 'master' into fix_xgboost_churn_neo

Author: Nikhil Kulkarni

sagemaker-experiments/mnist-handwritten-digits-classification-experiment/mnist-handwritten-digits-classification-experiment.ipynb

@@ -65,7 +65,7 @@
     "!{sys.executable} -m pip install torch==1.1.0\n",
     "!{sys.executable} -m pip install torchvision==0.3.0\n",
     "!{sys.executable} -m pip install pillow==6.2.2 ",
-    "!{sys.executable} -m pip install --upgrade sagemaker",
+    "!{sys.executable} -m pip install --upgrade sagemaker"
    ]
   },
   {

sagemaker_neo_compilation_jobs/pytorch_torchvision/code/resnet18.py

@@ -0,0 +1,78 @@
+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)
+
+
+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().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
+
+
+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, 'compiled.pt'))
+        model = model.to(device)
+
+        # It is recommended to run warm-up inference during model load
+        sample_input_path = os.path.join(model_dir, 'sample_input.pkl')
+        with open(sample_input_path, 'rb') as input_file:
+            model_input = pickle.load(input_file)
+        if torch.is_tensor(model_input):
+            model_input = model_input.to(device)
+            model(model_input)
+        elif isinstance(model_input, tuple):
+            model_input = (inp.to(device)
+                           for inp in model_input if torch.is_tensor(inp))
+            model(*model_input)
+        else:
+            print("Only supports a torch tensor or a tuple of torch tensors")
+
+        return model