From 14c1aae620f4e3160765c3006891704b4ecde5ba Mon Sep 17 00:00:00 2001 From: Gigon Bae Date: Fri, 8 Oct 2021 17:36:31 -0700 Subject: [PATCH] Force to remap to CPU for TorchScript model if GPU doesn't exist Signed-off-by: Gigon Bae --- .../mednist_classifier_monaideploy.py | 6 ++---- monai/deploy/core/models/torch_model.py | 4 +++- .../deploy/operators/monai_seg_inference_operator.py | 4 ++-- notebooks/tutorials/02_mednist_app-prebuilt.ipynb | 12 ++++-------- notebooks/tutorials/02_mednist_app.ipynb | 12 ++++-------- 5 files changed, 15 insertions(+), 23 deletions(-) diff --git a/examples/apps/mednist_classifier_monaideploy/mednist_classifier_monaideploy.py b/examples/apps/mednist_classifier_monaideploy/mednist_classifier_monaideploy.py index f180c73f..915d038d 100644 --- a/examples/apps/mednist_classifier_monaideploy/mednist_classifier_monaideploy.py +++ b/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() + 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) diff --git a/monai/deploy/core/models/torch_model.py b/monai/deploy/core/models/torch_model.py index bc390938..81a3f9ed 100644 --- a/monai/deploy/core/models/torch_model.py +++ b/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 diff --git a/monai/deploy/operators/monai_seg_inference_operator.py b/monai/deploy/operators/monai_seg_inference_operator.py index ec9d5c7c..df797e6a 100644 --- a/monai/deploy/operators/monai_seg_inference_operator.py +++ b/monai/deploy/operators/monai_seg_inference_operator.py @@ -158,6 +158,7 @@ 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. @@ -165,11 +166,10 @@ 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) + 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: diff --git a/notebooks/tutorials/02_mednist_app-prebuilt.ipynb b/notebooks/tutorials/02_mednist_app-prebuilt.ipynb index fe663a44..b0bf8457 100644 --- a/notebooks/tutorials/02_mednist_app-prebuilt.ipynb +++ b/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", diff --git a/notebooks/tutorials/02_mednist_app.ipynb b/notebooks/tutorials/02_mednist_app.ipynb index 16c54ea9..f128bf8b 100644 --- a/notebooks/tutorials/02_mednist_app.ipynb +++ b/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",