Add nb_epochs and bach_size to AdversarialTrainerMadryPGD.fit

Signed-off-by: Beat Buesser <[email protected]>

Author: Beat Buesser

art/defences/trainer/adversarial_trainer_madry_pgd.py

@@ -54,8 +54,8 @@ class AdversarialTrainerMadryPGD(Trainer):
     def __init__(
         self,
         classifier: "CLASSIFIER_LOSS_GRADIENTS_TYPE",
-        nb_epochs: int = 391,
-        batch_size: int = 128,
+        nb_epochs: Optional[int] = 391,
+        batch_size: Optional[int] = 128,
         eps: Union[int, float] = 8,
         eps_step: Union[int, float] = 2,
         max_iter: int = 7,
@@ -91,18 +91,36 @@ def __init__(
         self.trainer = AdversarialTrainer(classifier, self.attack, ratio=1.0)  # type: ignore
 
     def fit(  # pylint: disable=W0221
-        self, x: np.ndarray, y: np.ndarray, validation_data: Optional[np.ndarray] = None, **kwargs
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+        validation_data: Optional[np.ndarray] = None,
+        batch_size: Optional[int] = None,
+        nb_epochs: Optional[int] = None,
+        **kwargs
     ) -> None:
         """
         Train a model adversarially. See class documentation for more information on the exact procedure.
 
         :param x: Training data.
         :param y: Labels for the training data.
         :param validation_data: Validation data.
+        :param batch_size: Size of batches. Overwrites batch_size defined in __init__ if not None.
+        :param nb_epochs: Number of epochs to use for trainings. Overwrites nb_epochs defined in __init__ if not None.
         :param kwargs: Dictionary of framework-specific arguments.
         """
+        if batch_size is not None:
+            batch_size_fit = batch_size
+        else:
+            batch_size_fit = self.batch_size
+
+        if nb_epochs is not None:
+            nb_epochs_fit = nb_epochs
+        else:
+            nb_epochs_fit = self.nb_epochs
+
         self.trainer.fit(
-            x, y, validation_data=validation_data, nb_epochs=self.nb_epochs, batch_size=self.batch_size, **kwargs
+            x, y, validation_data=validation_data, nb_epochs=nb_epochs_fit, batch_size=batch_size_fit, **kwargs
         )
 
     def get_classifier(self) -> "CLASSIFIER_LOSS_GRADIENTS_TYPE":