Create new folder for any model name test

Author: saimidu

test/integration/__init__.py

@@ -27,6 +27,7 @@
 code_sub_dir = 'code'
 default_sub_dir = 'default_model'
 default_sub_eia_dir = 'default_model_eia'
+default_sub_traced_resnet_dir = 'default_traced_resnet'
 
 model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
 mnist_cpu_script = os.path.join(model_cpu_dir, code_sub_dir, 'mnist.py')
@@ -67,11 +68,13 @@
 default_model_tar = file_utils.make_tarfile(
     default_model_script, os.path.join(default_model_dir, "model.pt"), default_model_dir, script_path="code"
 )
+
+default_traced_resnet_dir = os.path.join(resnet18_path, default_sub_traced_resnet_dir)
+default_traced_resnet_script = os.path.join(default_traced_resnet_dir, code_sub_dir, "resnet18.py")
 default_model_traced_resnet18_tar = file_utils.make_tarfile(
-    default_model_script,
-    os.path.join(default_model_dir, "traced_resnet18.pt"),
-    default_model_dir,
-    filename="traced_resnet18.tar.gz",
+    default_traced_resnet_script,
+    os.path.join(default_traced_resnet_dir, "traced_resnet18.pt"),
+    default_traced_resnet_dir,
     script_path="code",
 )
 

test/integration/sagemaker/test_default_inference.py

@@ -23,6 +23,7 @@
 from integration import (
     default_model_script,
     default_model_tar,
+    default_traced_resnet_script,
     default_model_traced_resnet18_tar,
     default_model_eia_script,
     default_model_eia_tar,
@@ -74,11 +75,7 @@ def test_default_inference_any_model_name_gpu(sagemaker_session, image_uri, inst
     # Scripted model is serialized with torch.jit.save().
     # Default inference test doesn't need to instantiate model definition
     _test_default_inference(
-        sagemaker_session,
-        image_uri,
-        instance_type,
-        default_model_traced_resnet18_tar,
-        default_model_script,
+        sagemaker_session, image_uri, instance_type, default_model_traced_resnet18_tar, default_traced_resnet_script
     )
 
 

test/resources/resnet18/default_traced_resnet/code/resnet18.py

@@ -0,0 +1,51 @@
+import io
+import json
+import logging
+
+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 and 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