Update torch.export load with new api

examples/pt2/torch_export_aot_compile/install_pytorch_nightlies.sh

@@ -4,7 +4,11 @@
 pip uninstall torchtext torchdata torch torchvision torchaudio -y
 
 # Install nightly PyTorch and torchvision from the specified index URL
-pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 --ignore-installed
+if nvidia-smi > /dev/null 2>&1; then
+    pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 --ignore-installed
+else
+    pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu --ignore-installed
+fi
 
 # Optional: Display the installed PyTorch and torchvision versions
 python -c "import torch; print('PyTorch version:', torch.__version__)"

examples/pt2/torch_export_aot_compile/resnet18_torch_export.py

@@ -9,7 +9,15 @@
 model.eval()
 
 with torch.no_grad():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+        # The below config is needed for max batch_size = 16
+        # https://github.com/pytorch/pytorch/pull/116152
+        torch.backends.mkldnn.set_flags(False)
+        torch.backends.nnpack.set_flags(False)
+
     model = model.to(device=device)
     example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
 

test/pytest/test_torch_export.py

@@ -19,9 +19,9 @@
 MODEL_YAML_CFG_FILE = EXAMPLE_ROOT_DIR.joinpath("model-config.yaml")
 
 
-PT_220_AVAILABLE = (
+PT_230_AVAILABLE = (
     True
-    if packaging.version.parse(torch.__version__) > packaging.version.parse("2.1.1")
+    if packaging.version.parse(torch.__version__) > packaging.version.parse("2.2.2")
     else False
 )
 
@@ -47,7 +47,7 @@ def custom_working_directory(tmp_path):
     os.chdir(tmp_path)
 
 
-@pytest.mark.skipif(PT_220_AVAILABLE == False, reason="torch version is < 2.2.0")
+@pytest.mark.skipif(PT_230_AVAILABLE == False, reason="torch version is < 2.3.0")
 def test_torch_export_aot_compile(custom_working_directory):
     # Get the path to the custom working directory
     model_dir = custom_working_directory
@@ -88,7 +88,7 @@ def test_torch_export_aot_compile(custom_working_directory):
     assert labels == EXPECTED_RESULTS
 
 
-@pytest.mark.skipif(PT_220_AVAILABLE == False, reason="torch version is < 2.2.0")
+@pytest.mark.skipif(PT_230_AVAILABLE == False, reason="torch version is < 2.3.0")
 def test_torch_export_aot_compile_dynamic_batching(custom_working_directory):
     # Get the path to the custom working directory
     model_dir = custom_working_directory
@@ -123,6 +123,6 @@ def test_torch_export_aot_compile_dynamic_batching(custom_working_directory):
         data["body"] = byte_array_type
 
     # Send a batch of 16 elements
-    result = handler.handle([data for i in range(15)], ctx)
+    result = handler.handle([data for i in range(32)], ctx)
 
-    assert len(result) == 15
+    assert len(result) == 32

ts/torch_handler/base_handler.py

@@ -53,10 +53,10 @@
     )
     PT2_AVAILABLE = False
 
-if packaging.version.parse(torch.__version__) > packaging.version.parse("2.1.1"):
-    PT220_AVAILABLE = True
+if packaging.version.parse(torch.__version__) > packaging.version.parse("2.2.2"):
+    PT230_AVAILABLE = True
 else:
-    PT220_AVAILABLE = False
+    PT230_AVAILABLE = False
 
 if os.environ.get("TS_IPEX_ENABLE", "false") == "true":
     try:
@@ -187,7 +187,7 @@ def initialize(self, context):
         elif (
             self.model_pt_path.endswith(".so")
             and self._use_torch_export_aot_compile()
-            and PT220_AVAILABLE
+            and PT230_AVAILABLE
         ):
             # Set cuda device to the gpu_id of the backend worker
             # This is needed as the API for loading the exported model doesn't yet have a device id
@@ -256,9 +256,15 @@ def initialize(self, context):
         self.initialized = True
 
     def _load_torch_export_aot_compile(self, model_so_path):
-        from ts.handler_utils.torch_export.load_model import load_exported_model
+        """Loads the PyTorch model so and returns a Callable object.
 
-        return load_exported_model(model_so_path, self.map_location)
+        Args:
+            model_pt_path (str): denotes the path of the model file.
+
+        Returns:
+            (Callable Object) : Loads the model object.
+        """
+        return torch._export.aot_load(model_so_path, self.map_location)
 
     def _load_torchscript_model(self, model_pt_path):
         """Loads the PyTorch model and returns the NN model object.