Update Sagemaker Neo Notebooks examples to PyTorch 1.6 (#2018)

Co-authored-by: Prateek Chauhan <[email protected]>

Author: prateekch91

sagemaker_neo_compilation_jobs/pytorch_torchvision/code/resnet18.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import torch
+import neopytorch
 import torchvision.transforms as transforms
 from PIL import Image  # Training container doesn't have this package
 
@@ -55,24 +56,23 @@ def transform_fn(model, payload, request_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
+    neopytorch.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'), map_location=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

sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb

@@ -79,7 +79,7 @@
     "data_shape = '{\"input0\":[1,3,224,224]}'\n",
     "target_device = 'ml_c5'\n",
     "framework = 'PYTORCH'\n",
-    "framework_version = '1.4.0'\n",
+    "framework_version = '1.6'\n",
     "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)"
    ]
   },

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/code/vgg19_bn_compiled.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import torch
+import neopytorch
 import torchvision.transforms as transforms
 from PIL import Image  # Training container doesn't have this package
 
@@ -55,24 +56,23 @@ def transform_fn(model, payload, request_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)
+    neopytorch.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'), map_location=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")
+    # 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
+    return model

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/code/vgg19_bn_uncompiled.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import torch
+import neopytorch
 import torchvision.transforms as transforms
 from PIL import Image  # Training container doesn't have this package
 
@@ -55,10 +56,9 @@ def transform_fn(model, payload, request_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, 'model.pth'))
-        model = model.to(device)
+    neopytorch.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'), map_location=device)
 
-        return model
+    return model

sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb

@@ -24,7 +24,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.4.0 torchvision==0.5.0"
+    "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.6.0 torchvision==0.7.0"
    ]
   },
   {
@@ -127,7 +127,7 @@
     "data_shape = '{\"input0\":[1,3,224,224]}'\n",
     "target_device = 'ml_c5'\n",
     "framework = 'pytorch'\n",
-    "framework_version = '1.4.0'\n",
+    "framework_version = '1.6'\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)"