Force to remap to CPU for TorchScript model if GPUdoesn't exist

Signed-off-by: Gigon Bae <gbae@nvidia.com>

Author: gigony

examples/apps/mednist_classifier_monaideploy/mednist_classifier_monaideploy.py

@@ -66,12 +66,10 @@ def compute(self, op_input: InputContext, op_output: OutputContext, context: Exe
         image_tensor = self.transform(img)  # (1, 64, 64), torch.float64
         image_tensor = image_tensor[None].float()  # (1, 1, 64, 64), torch.float32
 
-        # Comment below line if you want to do CPU inference
-        image_tensor = image_tensor.cuda()
+        if torch.cuda.is_available():
+            image_tensor = image_tensor.cuda()
 
         model = context.models.get()  # get a TorchScriptModel object
-        # Uncomment the following line if you want to do CPU inference
-        # model.predictor = torch.jit.load(model.path, map_location="cpu").eval()
 
         with torch.no_grad():
             outputs = model(image_tensor)

monai/deploy/core/models/torch_model.py

@@ -47,7 +47,9 @@ def predictor(self) -> "torch.nn.Module":  # type: ignore
             torch.nn.Module: the model's predictor
         """
         if self._predictor is None:
-            self._predictor = torch.jit.load(self.path).eval()
+            # If GPU is not available, explicitly use "cpu" to dynamically remap storages to CPU.
+            map_location = None if torch.cuda.is_available() else "cpu"
+            self._predictor = torch.jit.load(self.path, map_location=map_location).eval()
         return self._predictor
 
     @predictor.setter

monai/deploy/operators/monai_seg_inference_operator.py

@@ -165,7 +165,8 @@ def compute(self, op_input: InputContext, op_output: OutputContext, context: Exe
                 model = context.models.get()
             else:
                 print(f"Loading TorchScript model from: {MonaiSegInferenceOperator.MODEL_LOCAL_PATH}")
-                model = torch.jit.load(MonaiSegInferenceOperator.MODEL_LOCAL_PATH)
+                map_location = None if torch.cuda.is_available() else "cpu"
+                model = torch.jit.load(MonaiSegInferenceOperator.MODEL_LOCAL_PATH, map_location=map_location)
 
             dataset = Dataset(data=[{self._input_dataset_key: img_name}], transform=pre_transforms)
             dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)

notebooks/tutorials/02_mednist_app-prebuilt.ipynb

@@ -432,12 +432,10 @@
     "        image_tensor = self.transform(img)  # (1, 64, 64), torch.float64\n",
     "        image_tensor = image_tensor[None].float()  # (1, 1, 64, 64), torch.float32\n",
     "\n",
-    "        # Comment below line if you want to do CPU inference\n",
-    "        image_tensor = image_tensor.cuda()\n",
+    "        if torch.cuda.is_available():\n",
+    "            image_tensor = image_tensor.cuda()\n",
     "\n",
     "        model = context.models.get()  # get a TorchScriptModel object\n",
-    "        # Uncomment the following line if you want to do CPU inference\n",
-    "        # model.predictor = torch.jit.load(model.path, map_location=\"cpu\").eval()\n",
     "\n",
     "        with torch.no_grad():\n",
     "            outputs = model(image_tensor)\n",
@@ -652,12 +650,10 @@
     "        image_tensor = self.transform(img)  # (1, 64, 64), torch.float64\n",
     "        image_tensor = image_tensor[None].float()  # (1, 1, 64, 64), torch.float32\n",
     "\n",
-    "        # Comment below line if you want to do CPU inference\n",
-    "        image_tensor = image_tensor.cuda()\n",
+    "        if torch.cuda.is_available():\n",
+    "            image_tensor = image_tensor.cuda()\n",
     "\n",
     "        model = context.models.get()  # get a TorchScriptModel object\n",
-    "        # Uncomment the following line if you want to do CPU inference\n",
-    "        # model.predictor = torch.jit.load(model.path, map_location=\"cpu\").eval()\n",
     "\n",
     "        with torch.no_grad():\n",
     "            outputs = model(image_tensor)\n",

notebooks/tutorials/02_mednist_app.ipynb

@@ -510,12 +510,10 @@
     "        image_tensor = self.transform(img)  # (1, 64, 64), torch.float64\n",
     "        image_tensor = image_tensor[None].float()  # (1, 1, 64, 64), torch.float32\n",
     "\n",
-    "        # Comment below line if you want to do CPU inference\n",
-    "        image_tensor = image_tensor.cuda()\n",
+    "        if torch.cuda.is_available():\n",
+    "            image_tensor = image_tensor.cuda()\n",
     "\n",
     "        model = context.models.get()  # get a TorchScriptModel object\n",
-    "        # Uncomment the following line if you want to do CPU inference\n",
-    "        # model.predictor = torch.jit.load(model.path, map_location=\"cpu\").eval()\n",
     "\n",
     "        with torch.no_grad():\n",
     "            outputs = model(image_tensor)\n",
@@ -748,12 +746,10 @@
     "        image_tensor = self.transform(img)  # (1, 64, 64), torch.float64\n",
     "        image_tensor = image_tensor[None].float()  # (1, 1, 64, 64), torch.float32\n",
     "\n",
-    "        # Comment below line if you want to do CPU inference\n",
-    "        image_tensor = image_tensor.cuda()\n",
+    "        if torch.cuda.is_available():\n",
+    "            image_tensor = image_tensor.cuda()\n",
     "\n",
     "        model = context.models.get()  # get a TorchScriptModel object\n",
-    "        # Uncomment the following line if you want to do CPU inference\n",
-    "        # model.predictor = torch.jit.load(model.path, map_location=\"cpu\").eval()\n",
     "\n",
     "        with torch.no_grad():\n",
     "            outputs = model(image_tensor)\n",