Introduce fp16 flag to enable/disable mixed precision for predict()

src/super_gradients/training/models/classification_models/base_classifer.py

@@ -30,10 +30,11 @@ def set_dataset_processing_params(self, class_names: Optional[List[str]] = None,
         self._image_processor = image_processor or self._image_processor
 
     @lru_cache(maxsize=1)
-    def _get_pipeline(self, fuse_model: bool = True, skip_image_resizing: bool = False) -> ClassificationPipeline:
+    def _get_pipeline(self, fuse_model: bool = True, skip_image_resizing: bool = False, fp16: bool = True) -> ClassificationPipeline:
         """Instantiate the prediction pipeline of this model.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16: If True, use mixed precision for inference.
         """
         if None in (self._class_names, self._image_processor):
             raise RuntimeError(
@@ -48,24 +49,34 @@ def _get_pipeline(self, fuse_model: bool = True, skip_image_resizing: bool = Fal
             image_processor=self._image_processor,
             class_names=self._class_names,
             fuse_model=fuse_model,
+            fp16=fp16,
         )
         return pipeline
 
-    def predict(self, images: ImageSource, batch_size: int = 32, fuse_model: bool = True, skip_image_resizing: bool = False) -> ImagesClassificationPrediction:
+    def predict(
+        self,
+        images: ImageSource,
+        batch_size: int = 32,
+        fuse_model: bool = True,
+        skip_image_resizing: bool = False,
+        fp16: bool = True,
+    ) -> ImagesClassificationPrediction:
         """Predict an image or a list of images.
 
         :param images:      Images to predict.
         :param batch_size:  Maximum number of images to process at the same time.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16: If True, use mixed precision for inference.
         """
-        pipeline = self._get_pipeline(fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         return pipeline(images, batch_size=batch_size)  # type: ignore
 
-    def predict_webcam(self, fuse_model: bool = True, skip_image_resizing: bool = False) -> None:
+    def predict_webcam(self, fuse_model: bool = True, skip_image_resizing: bool = False, fp16: bool = True) -> None:
         """Predict using webcam.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16: If True, use mixed precision for inference.
         """
-        pipeline = self._get_pipeline(fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         pipeline.predict_webcam()

src/super_gradients/training/models/detection_models/customizable_detector.py

@@ -5,6 +5,7 @@
  * each module accepts in_channels and other parameters
  * each module defines out_channels property on construction
 """
+
 from typing import Union, Optional, List, Callable
 from functools import lru_cache
 
@@ -208,6 +209,7 @@ def _get_pipeline(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ) -> DetectionPipeline:
         """Instantiate the prediction pipeline of this model.
 
@@ -222,6 +224,7 @@ def _get_pipeline(
                                     If False, each anchor can produce only one label of the class with the highest score.
         :param class_agnostic_nms:  (Optional) If True, perform class-agnostic NMS (i.e IoU of boxes of different classes is checked).
                                     If False NMS is performed separately for each class.
+        :param fp16:                If True, use mixed precision for inference.
         """
         if None in (self._class_names, self._image_processor, self._default_nms_iou, self._default_nms_conf):
             raise RuntimeError(
@@ -256,6 +259,7 @@ def _get_pipeline(
             ),
             class_names=self._class_names,
             fuse_model=fuse_model,
+            fp16=fp16,
         )
         return pipeline
 
@@ -271,6 +275,7 @@ def predict(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ) -> ImagesDetectionPrediction:
         """Predict an image or a list of images.
 
@@ -287,6 +292,7 @@ def predict(
                                     If False, each anchor can produce only one label of the class with the highest score.
         :param class_agnostic_nms:  (Optional) If True, perform class-agnostic NMS (i.e IoU of boxes of different classes is checked).
                                     If False NMS is performed separately for each class.
+        :param fp16:                        If True, use mixed precision for inference.
         """
         pipeline = self._get_pipeline(
             iou=iou,
@@ -297,6 +303,7 @@ def predict(
             max_predictions=max_predictions,
             multi_label_per_box=multi_label_per_box,
             class_agnostic_nms=class_agnostic_nms,
+            fp16=fp16,
         )
         return pipeline(images, batch_size=batch_size)  # type: ignore
 
@@ -310,6 +317,7 @@ def predict_webcam(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ):
         """Predict using webcam.
 
@@ -325,6 +333,7 @@ def predict_webcam(
                                     If False, each anchor can produce only one label of the class with the highest score.
         :param class_agnostic_nms:  (Optional) If True, perform class-agnostic NMS (i.e IoU of boxes of different classes is checked).
                                     If False NMS is performed separately for each class.
+        :param fp16:                If True, use mixed precision for inference.
         """
         pipeline = self._get_pipeline(
             iou=iou,
@@ -335,6 +344,7 @@ def predict_webcam(
             max_predictions=max_predictions,
             multi_label_per_box=multi_label_per_box,
             class_agnostic_nms=class_agnostic_nms,
+            fp16=fp16,
         )
         pipeline.predict_webcam()
 

src/super_gradients/training/models/detection_models/pp_yolo_e/pp_yolo_e.py

@@ -208,6 +208,7 @@ def _get_pipeline(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ) -> DetectionPipeline:
         """Instantiate the prediction pipeline of this model.
 
@@ -256,6 +257,7 @@ def _get_pipeline(
             ),
             class_names=self._class_names,
             fuse_model=fuse_model,
+            fp16=fp16,
         )
         return pipeline
 
@@ -271,6 +273,7 @@ def predict(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ) -> ImagesDetectionPrediction:
         """Predict an image or a list of images.
 
@@ -287,6 +290,7 @@ def predict(
                                     If False, each anchor can produce only one label of the class with the highest score.
         :param class_agnostic_nms:  (Optional) If True, perform class-agnostic NMS (i.e IoU of boxes of different classes is checked).
                                     If False NMS is performed separately for each class.
+        :param fp16:                If True, the model will use mixed precision for inference.
         """
         pipeline = self._get_pipeline(
             iou=iou,
@@ -297,6 +301,7 @@ def predict(
             max_predictions=max_predictions,
             multi_label_per_box=multi_label_per_box,
             class_agnostic_nms=class_agnostic_nms,
+            fp16=fp16,
         )
         return pipeline(images, batch_size=batch_size)  # type: ignore
 
@@ -310,13 +315,13 @@ def predict_webcam(
         max_predictions: Optional[int] = None,
         multi_label_per_box: Optional[bool] = None,
         class_agnostic_nms: Optional[bool] = None,
+        fp16: bool = True,
     ):
         """Predict using webcam.
 
         :param iou:                 (Optional) IoU threshold for the nms algorithm. If None, the default value associated to the training is used.
         :param conf:                (Optional) Below the confidence threshold, prediction are discarded.
                                     If None, the default value associated to the training is used.
-        :param batch_size:          Maximum number of images to process at the same time.
         :param fuse_model:          If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
         :param nms_top_k:           (Optional) The maximum number of detections to consider for NMS.
@@ -325,6 +330,7 @@ def predict_webcam(
                                     If False, each anchor can produce only one label of the class with the highest score.
         :param class_agnostic_nms:  (Optional) If True, perform class-agnostic NMS (i.e IoU of boxes of different classes is checked).
                                     If False NMS is performed separately for each class.
+        :param fp16:                If True, use mixed precision for inference.
         """
         pipeline = self._get_pipeline(
             iou=iou,
@@ -335,6 +341,7 @@ def predict_webcam(
             max_predictions=max_predictions,
             multi_label_per_box=multi_label_per_box,
             class_agnostic_nms=class_agnostic_nms,
+            fp16=fp16,
         )
         pipeline.predict_webcam()
 

src/super_gradients/training/models/detection_models/yolo_base.py

@@ -537,7 +537,7 @@ def set_dataset_processing_params(
 
     @lru_cache(maxsize=1)
     def _get_pipeline(
-        self, iou: Optional[float] = None, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False
+        self, iou: Optional[float] = None, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False, fp16: bool = True
     ) -> DetectionPipeline:
         """Instantiate the prediction pipeline of this model.
 
@@ -546,6 +546,7 @@ def _get_pipeline(
                         If None, the default value associated to the training is used.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16:    If True, use mixed precision for inference.
         """
         if None in (self._class_names, self._image_processor, self._default_nms_iou, self._default_nms_conf):
             raise RuntimeError(
@@ -569,6 +570,7 @@ def _get_pipeline(
             post_prediction_callback=self.get_post_prediction_callback(iou=iou, conf=conf),
             class_names=self._class_names,
             fuse_model=fuse_model,
+            fp16=fp16,
         )
         return pipeline
 
@@ -580,6 +582,7 @@ def predict(
         batch_size: int = 32,
         fuse_model: bool = True,
         skip_image_resizing: bool = False,
+        fp16: bool = True,
     ) -> ImagesDetectionPrediction:
         """Predict an image or a list of images.
 
@@ -590,20 +593,22 @@ def predict(
         :param batch_size:  Maximum number of images to process at the same time.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16:        If True, use mixed precision for inference.
         """
-        pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         return pipeline(images, batch_size=batch_size)  # type: ignore
 
-    def predict_webcam(self, iou: Optional[float] = None, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False):
+    def predict_webcam(self, iou: Optional[float] = None, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False, fp16=True):
         """Predict using webcam.
 
         :param iou:     (Optional) IoU threshold for the nms algorithm. If None, the default value associated to the training is used.
         :param conf:    (Optional) Below the confidence threshold, prediction are discarded.
                         If None, the default value associated to the training is used.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16:    If True, use mixed precision for inference.
         """
-        pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         pipeline.predict_webcam()
 
     def train(self, mode: bool = True):

src/super_gradients/training/models/pose_estimation_models/dekr_hrnet.py

@@ -583,7 +583,9 @@ def set_dataset_processing_params(
         self._default_nms_conf = conf or self._default_nms_conf
 
     @lru_cache(maxsize=1)
-    def _get_pipeline(self, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False) -> PoseEstimationPipeline:
+    def _get_pipeline(
+        self, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False, fp16: bool = True
+    ) -> PoseEstimationPipeline:
         """Instantiate the prediction pipeline of this model.
 
         :param conf:    (Optional) Below the confidence threshold, prediction are discarded.
@@ -621,11 +623,18 @@ def _get_pipeline(self, conf: Optional[float] = None, fuse_model: bool = True, s
             keypoint_colors=self._keypoint_colors,
             post_prediction_callback=self.get_post_prediction_callback(conf=conf),
             fuse_model=fuse_model,
+            fp16=fp16,
         )
         return pipeline
 
     def predict(
-        self, images: ImageSource, conf: Optional[float] = None, batch_size: int = 32, fuse_model: bool = True, skip_image_resizing: bool = False
+        self,
+        images: ImageSource,
+        conf: Optional[float] = None,
+        batch_size: int = 32,
+        fuse_model: bool = True,
+        skip_image_resizing: bool = False,
+        fp16: bool = True,
     ) -> ImagesPoseEstimationPrediction:
         """Predict an image or a list of images.
 
@@ -635,19 +644,20 @@ def predict(
         :param batch_size:  Maximum number of images to process at the same time.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
+        :param fp16:       If True, use mixed precision for inference.
         """
-        pipeline = self._get_pipeline(conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         return pipeline(images, batch_size=batch_size)  # type: ignore
 
-    def predict_webcam(self, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False):
+    def predict_webcam(self, conf: Optional[float] = None, fuse_model: bool = True, skip_image_resizing: bool = False, fp16: bool = True):
         """Predict using webcam.
 
         :param conf:    (Optional) Below the confidence threshold, prediction are discarded.
                         If None, the default value associated to the training is used.
         :param fuse_model:  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
         :param skip_image_resizing: If True, the image processor will not resize the images.
         """
-        pipeline = self._get_pipeline(conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing)
+        pipeline = self._get_pipeline(conf=conf, fuse_model=fuse_model, skip_image_resizing=skip_image_resizing, fp16=fp16)
         pipeline.predict_webcam()
 
     def train(self, mode: bool = True):