Add a flag to be used for marking the YOLOv8 model

Author: CNOCycle

art/estimators/object_detection/pytorch_object_detector.py

@@ -65,6 +65,7 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
+        is_yolov8: bool = False,
     ):
         """
         Initialization.
@@ -92,6 +93,7 @@ def __init__(
                               'loss_objectness', and 'loss_rpn_box_reg'.
         :param device_type: Type of device to be used for model and tensors, if `cpu` run on CPU, if `gpu` run on GPU
                             if available otherwise run on CPU.
+        :param is_yolov8: The flag to be used for marking the YOLOv8 model.
         """
         import re
         import torch
@@ -140,9 +142,10 @@ def __init__(
 
         self._model: torch.nn.Module
         self._model.to(self._device)
-        try:
+        self.is_yolov8 = is_yolov8
+        if self.is_yolov8:
             self._model.model.eval()
-        except AttributeError:
+        else:
             self._model.eval()
 
     @property
@@ -406,9 +409,9 @@ def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> list[dict[s
         from torch.utils.data import TensorDataset, DataLoader
 
         # Set model to evaluation mode
-        try:
+        if self.is_yolov8:
             self._model.model.eval()
-        except AttributeError:
+        else:
             self._model.eval()
 
         # Apply preprocessing and convert to tensors

art/estimators/object_detection/pytorch_yolo.py

@@ -64,6 +64,7 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
+        is_yolov8: bool = False,
     ):
         """
         Initialization.
@@ -92,6 +93,7 @@ def __init__(
                               'loss_objectness', and 'loss_rpn_box_reg'.
         :param device_type: Type of device to be used for model and tensors, if `cpu` run on CPU, if `gpu` run on GPU
                             if available otherwise run on CPU.
+        :param is_yolov8: The flag to be used for marking the YOLOv8 model.
         """
         super().__init__(
             model=model,
@@ -104,6 +106,7 @@ def __init__(
             preprocessing=preprocessing,
             attack_losses=attack_losses,
             device_type=device_type,
+            is_yolov8=is_yolov8,
         )
 
     def _translate_labels(self, labels: list[dict[str, "torch.Tensor"]]) -> "torch.Tensor":

notebooks/snal.ipynb

@@ -71,7 +71,8 @@
     "model = YOLO('yolov8m')\n",
     "py_model = PyTorchYolo(model=model,\n",
     "                       input_shape=(3, 640, 640),\n",
-    "                       channels_first=True)\n",
+    "                       channels_first=True,\n",
+    "                       is_yolov8=True)\n",
     "\n",
     "# Define a custom function to collect patches from images\n",
     "def collect_patches_from_images(model: \"torch.nn.Module\",\n",

tests/attacks/evasion/test_steal_now_attack_later.py

@@ -35,7 +35,7 @@ def test_generate(art_warning):
         import requests
 
         model = YOLO("yolov8m")
-        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)
+        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)
 
         # Define a custom function to collect patches from images
         def collect_patches_from_images(model, imgs):
@@ -192,7 +192,7 @@ def _loader(self, path):
 def test_check_params(art_warning):
     try:
         model = YOLO("yolov8m")
-        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)
+        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)
 
         def dummy_func(model, imags):
             candidates_patch = []