🐞 Add device flag

anomalib/deploy/inferencers/torch_inferencer.py

@@ -30,15 +30,19 @@ class TorchInferencer(Inferencer):
         model_source (Union[str, Path, AnomalyModule]): Path to the model ckpt file or the Anomaly model.
         meta_data_path (Union[str, Path], optional): Path to metadata file. If none, it tries to load the params
                 from the model state_dict. Defaults to None.
+        device (Optional[str], optional): Device to use for inference. Options are auto, cpu, cuda. Defaults to "auto".
     """
 
     def __init__(
         self,
         config: Union[str, Path, DictConfig, ListConfig],
         model_source: Union[str, Path, AnomalyModule],
         meta_data_path: Union[str, Path] = None,
+        device: str = "auto",
     ):
 
+        self.device = self._get_device(device)
+
         # Check and load the configuration
         if isinstance(config, (str, Path)):
             self.config = get_configurable_parameters(config_path=config)
@@ -55,6 +59,22 @@ def __init__(
 
         self.meta_data = self._load_meta_data(meta_data_path)
 
+    def _get_device(self, device: str) -> torch.device:
+        """Get the device to use for inference.
+
+        Args:
+            device (str): Device to use for inference. Options are auto, cpu, cuda.
+
+        Returns:
+            torch.device: Device to use for inference.
+        """
+        if device not in ("auto", "cpu", "cuda"):
+            raise ValueError(f"Unknown device {device}")
+
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        return torch.device(device)
+
     def _load_meta_data(self, path: Optional[Union[str, Path]] = None) -> Union[Dict, DictConfig]:
         """Load metadata from file or from model state dict.
 
@@ -84,7 +104,7 @@ def load_model(self, path: Union[str, Path]) -> AnomalyModule:
         model = get_model(self.config)
         model.load_state_dict(torch.load(path)["state_dict"])
         model.eval()
-        return model
+        return model.to(self.device)
 
     def pre_process(self, image: np.ndarray) -> Tensor:
         """Pre process the input image by applying transformations.
@@ -105,7 +125,7 @@ def pre_process(self, image: np.ndarray) -> Tensor:
         if len(processed_image) == 3:
             processed_image = processed_image.unsqueeze(0)
 
-        return processed_image
+        return processed_image.to(self.device)
 
     def forward(self, image: Tensor) -> Tensor:
         """Forward-Pass input tensor to the model.

tools/inference/torch_inference.py

@@ -31,6 +31,13 @@ def get_args() -> Namespace:
     parser.add_argument("--weights", type=Path, required=True, help="Path to model weights")
     parser.add_argument("--input", type=Path, required=True, help="Path to an image to infer.")
     parser.add_argument("--output", type=Path, required=False, help="Path to save the output image.")
+    parser.add_argument(
+        "--device",
+        type=str,
+        required=False,
+        default="auto",
+        help="Device to use for inference. Defaults to auto. Options: auto, cpu, cuda",
+    )
     parser.add_argument(
         "--task",
         type=str,
@@ -70,7 +77,7 @@ def infer() -> None:
     args = get_args()
 
     # Create the inferencer and visualizer.
-    inferencer = TorchInferencer(config=args.config, model_source=args.weights)
+    inferencer = TorchInferencer(config=args.config, model_source=args.weights, device=args.device)
     visualizer = Visualizer(mode=args.visualization_mode, task=args.task)
 
     filenames = get_image_filenames(path=args.input)