fix kwargs bug

Signed-off-by: GiulioZizzo <[email protected]>

Author: GiulioZizzo

art/estimators/certification/derandomized_smoothing/derandomized_smoothing.py

@@ -91,12 +91,13 @@ def _predict_classifier(self, x: np.ndarray, batch_size: int, training_mode: boo
         """
         raise NotImplementedError
 
-    def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:
+    def predict(self, x: np.ndarray, batch_size: int = 128, training_mode: bool = False, **kwargs) -> np.ndarray:
         """
         Performs cumulative predictions over every ablation location
 
         :param x: Unablated image
         :param batch_size: the batch size for the prediction
+        :param training_mode: if to run the classifier in training mode
         :return: cumulative predictions after sweeping over all the ablation configurations.
         """
         if self._channels_first:
@@ -116,20 +117,24 @@ def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:
             for ablation_start in range(ablate_over_range):
                 ablated_x = self.ablator.forward(np.copy(x), column_pos=ablation_start)
                 if ablation_start == 0:
-                    preds = self._predict_classifier(ablated_x, batch_size=batch_size, training_mode=False, **kwargs)
+                    preds = self._predict_classifier(
+                        ablated_x, batch_size=batch_size, training_mode=training_mode, **kwargs
+                    )
                 else:
-                    preds += self._predict_classifier(ablated_x, batch_size=batch_size, training_mode=False, **kwargs)
+                    preds += self._predict_classifier(
+                        ablated_x, batch_size=batch_size, training_mode=training_mode, **kwargs
+                    )
         elif self.ablation_type == "block":
             for xcorner in range(rows_in_data):
                 for ycorner in range(columns_in_data):
                     ablated_x = self.ablator.forward(np.copy(x), row_pos=xcorner, column_pos=ycorner)
                     if ycorner == 0 and xcorner == 0:
                         preds = self._predict_classifier(
-                            ablated_x, batch_size=batch_size, training_mode=False, **kwargs
+                            ablated_x, batch_size=batch_size, training_mode=training_mode, **kwargs
                         )
                     else:
                         preds += self._predict_classifier(
-                            ablated_x, batch_size=batch_size, training_mode=False, **kwargs
+                            ablated_x, batch_size=batch_size, training_mode=training_mode, **kwargs
                         )
         return preds
 

art/estimators/certification/derandomized_smoothing/pytorch.py

@@ -131,15 +131,18 @@ def _predict_classifier(self, x: np.ndarray, batch_size: int, training_mode: boo
             (torch.nn.functional.softmax(torch.from_numpy(outputs), dim=1) >= self.threshold).type(torch.int)
         )
 
-    def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:  # type: ignore
+    def predict(
+        self, x: np.ndarray, batch_size: int = 128, training_mode: bool = False, **kwargs
+    ) -> np.ndarray:  # type: ignore
         """
         Perform prediction of the given classifier for a batch of inputs, taking an expectation over transformations.
 
         :param x: Input samples.
         :param batch_size: Batch size.
+        :param training_mode: if to run the classifier in training mode
         :return: Array of predictions of shape `(nb_inputs, nb_classes)`.
         """
-        return DeRandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=False, **kwargs)
+        return DeRandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=training_mode, **kwargs)
 
     def _fit_classifier(self, x: np.ndarray, y: np.ndarray, batch_size: int, nb_epochs: int, **kwargs) -> None:
         x = x.astype(ART_NUMPY_DTYPE)

art/estimators/certification/derandomized_smoothing/tensorflow.py

@@ -170,12 +170,15 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
             if scheduler is not None:
                 scheduler(epoch)
 
-    def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:  # type: ignore
+    def predict(
+        self, x: np.ndarray, batch_size: int = 128, training_mode: bool = False, **kwargs
+    ) -> np.ndarray:  # type: ignore
         """
         Perform prediction of the given classifier for a batch of inputs
 
         :param x: Input samples.
         :param batch_size: Batch size.
+        :param training_mode: if to run the classifier in training mode
         :return: Array of predictions of shape `(nb_inputs, nb_classes)`.
         """
-        return DeRandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=False, **kwargs)
+        return DeRandomizedSmoothingMixin.predict(self, x, batch_size=batch_size, training_mode=training_mode, **kwargs)