Merge pull request #167 from Project-MONAI/force_cpu_torchscript

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

Author: gigony

examples/apps/mednist_classifier_monaideploy/mednist_classifier_monaideploy.py

@@ -69,12 +69,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()
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        image_tensor = image_tensor.to(device)
 
         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()
+            # Use a device to dynamically remap, depending on the GPU availability.
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self._predictor = torch.jit.load(self.path, map_location=device).eval()
         return self._predictor
 
     @predictor.setter

monai/deploy/operators/monai_seg_inference_operator.py

@@ -195,18 +195,18 @@ def compute(self, op_input: InputContext, op_output: OutputContext, context: Exe
             pre_transforms = self._pre_transform if self._pre_transform else self.pre_process(self._reader)
             post_transforms = self._post_transforms if self._post_transforms else self.post_process(pre_transforms)
 
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             model = None
             if context.models:
                 # `context.models.get(model_name)` returns a model instance if exists.
                 # If model_name is not specified and only one model exists, it returns that model.
                 model = context.models.get()
             else:
                 print(f"Loading TorchScript model from: {MonaiSegInferenceOperator.MODEL_LOCAL_PATH}")
-                model = torch.jit.load(MonaiSegInferenceOperator.MODEL_LOCAL_PATH)
+                model = torch.jit.load(MonaiSegInferenceOperator.MODEL_LOCAL_PATH, map_location=device)
 
             dataset = Dataset(data=[{self._input_dataset_key: img_name}], transform=pre_transforms)
             dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
             with torch.no_grad():
                 for d in dataloader:

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",
+    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "        image_tensor = image_tensor.to(device)\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",
+    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "        image_tensor = image_tensor.to(device)\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",
+    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "        image_tensor = image_tensor.to(device)\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",
+    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "        image_tensor = image_tensor.to(device)\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",