standardise verbose use across tools

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

Author: GiulioZizzo

art/estimators/certification/derandomized_smoothing/pytorch.py

@@ -441,7 +441,7 @@ def fit(  # pylint: disable=W0221
         update_batchnorm: bool = True,
         batchnorm_update_epochs: int = 1,
         transform: Optional["torchvision.transforms.transforms.Compose"] = None,
-        verbose: bool = True,
+        verbose: Optional[Union[bool, int]] = None,
         **kwargs,
     ) -> None:
         """
@@ -469,6 +469,17 @@ def fit(  # pylint: disable=W0221
         """
         import torch
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = False
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -499,7 +510,7 @@ def fit(  # pylint: disable=W0221
             epoch_loss = []
             epoch_batch_sizes = []
 
-            pbar = tqdm(range(num_batch), disable=not verbose)
+            pbar = tqdm(range(num_batch), disable=not display_pb)
 
             # Train for one epoch
             for m in pbar:
@@ -545,7 +556,7 @@ def fit(  # pylint: disable=W0221
                 epoch_loss.append(loss.cpu().detach().numpy())
                 epoch_batch_sizes.append(len(i_batch))
 
-                if verbose:
+                if display_pb:
                     pbar.set_description(
                         f"Loss {np.average(epoch_loss, weights=epoch_batch_sizes):.3f} "
                         f"Acc {np.average(epoch_acc, weights=epoch_batch_sizes):.3f} "

art/estimators/certification/derandomized_smoothing/tensorflow.py

@@ -155,7 +155,13 @@ def _predict_classifier(self, x: np.ndarray, batch_size: int, training_mode: boo
         return np.asarray(outputs >= self.threshold).astype(int)
 
     def fit(  # pylint: disable=W0221
-        self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 10, verbose: bool = True, **kwargs
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+        batch_size: int = 128,
+        nb_epochs: int = 10,
+        verbose: Optional[Union[bool, int]] = None,
+        **kwargs,
     ) -> None:
         """
         Fit the classifier on the training set `(x, y)`.
@@ -165,13 +171,24 @@ def fit(  # pylint: disable=W0221
                   shape (nb_samples,).
         :param batch_size: Size of batches.
         :param nb_epochs: Number of epochs to use for training.
-        :param verbose: if to display training progress bars
+        :param verbose: If to display training progress bars
         :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
                        "scheduler" which is an optional function that will be called at the end of every
                        epoch to adjust the learning rate.
         """
         import tensorflow as tf
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = False
+
         if self._train_step is None:  # pragma: no cover
             if self._loss_object is None:  # pragma: no cover
                 raise TypeError(
@@ -214,7 +231,7 @@ def train_step(model, images, labels):
             epoch_loss = []
             epoch_batch_sizes = []
 
-            pbar = tqdm(range(num_batch), disable=not verbose)
+            pbar = tqdm(range(num_batch), disable=not display_pb)
 
             ind = np.arange(len(x_preprocessed))
             for m in pbar:
@@ -231,7 +248,7 @@ def train_step(model, images, labels):
                 else:
                     train_step(self.model, images, labels)
 
-                if verbose:
+                if display_pb:
                     if self._train_step is None:
                         pbar.set_description(
                             f"Loss {np.average(epoch_loss, weights=epoch_batch_sizes):.3f} "

art/estimators/certification/randomized_smoothing/macer/pytorch.py

@@ -138,6 +138,7 @@ def fit(  # pylint: disable=W0221
         training_mode: bool = True,
         drop_last: bool = False,
         scheduler: Optional["torch.optim.lr_scheduler._LRScheduler"] = None,
+        verbose: Optional[Union[bool, int]] = None,
         **kwargs,
     ) -> None:
         """
@@ -153,13 +154,26 @@ def fit(  # pylint: disable=W0221
                           the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                           the last batch will be smaller. (default: ``False``)
         :param scheduler: Learning rate scheduler to run at the start of every epoch.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                        class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                and providing it takes no effect.
         """
         import torch
         import torch.nn.functional as F
         from torch.utils.data import TensorDataset, DataLoader
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -185,7 +199,7 @@ def fit(  # pylint: disable=W0221
         )
 
         # Start training
-        for _ in trange(nb_epochs, disable=not self.verbose):
+        for _ in trange(nb_epochs, disable=not display_pb):
             for x_batch, y_batch in dataloader:
                 # Move inputs to GPU
                 x_batch = x_batch.to(self.device)

art/estimators/certification/randomized_smoothing/macer/tensorflow.py

@@ -132,7 +132,15 @@ def __init__(
         self.lmbda = lmbda
         self.gaussian_samples = gaussian_samples
 
-    def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 10, **kwargs) -> None:
+    def fit(
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+        batch_size: int = 128,
+        nb_epochs: int = 10,
+        verbose: Optional[Union[bool, int]] = None,
+        **kwargs
+    ) -> None:
         """
         Fit the classifier on the training set `(x, y)`.
 
@@ -141,12 +149,25 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                   shape (nb_samples,).
         :param batch_size: Size of batches.
         :param nb_epochs: Number of epochs to use for training.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                        class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
                        "scheduler" which is an optional function that will be called at the end of every
                        epoch to adjust the learning rate.
         """
         import tensorflow as tf
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         if self._train_step is None:  # pragma: no cover
             if self._optimizer is None:  # pragma: no cover
                 raise ValueError(
@@ -213,7 +234,7 @@ def train_step(model, images, labels):
 
         train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)
 
-        for epoch in trange(nb_epochs, disable=not self.verbose):
+        for epoch in trange(nb_epochs, disable=not display_pb):
             for images, labels in train_ds:
                 # Tile samples for Gaussian augmentation
                 input_size = len(images)

art/estimators/certification/randomized_smoothing/pytorch.py

@@ -140,6 +140,7 @@ def fit(  # pylint: disable=W0221
         training_mode: bool = True,
         drop_last: bool = False,
         scheduler: Optional["torch.optim.lr_scheduler._LRScheduler"] = None,
+        verbose: Optional[Union[bool, int]] = None,
         **kwargs,
     ) -> None:
         """
@@ -155,12 +156,25 @@ def fit(  # pylint: disable=W0221
                           the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                           the last batch will be smaller. (default: ``False``)
         :param scheduler: Learning rate scheduler to run at the start of every epoch.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                 class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                and providing it takes no effect.
         """
         import torch
         from torch.utils.data import TensorDataset, DataLoader
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -182,7 +196,7 @@ def fit(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
 
         # Start training
-        for _ in trange(nb_epochs, disable=not self.verbose):
+        for _ in trange(nb_epochs, disable=not display_pb):
             for x_batch, y_batch in dataloader:
                 # Move inputs to device
                 x_batch = x_batch.to(self._device)

art/estimators/certification/randomized_smoothing/smooth_adv/pytorch.py

@@ -155,6 +155,7 @@ def fit(  # pylint: disable=W0221
         training_mode: bool = True,
         drop_last: bool = False,
         scheduler: Optional["torch.optim.lr_scheduler._LRScheduler"] = None,
+        verbose: Optional[Union[bool, int]] = None,
         **kwargs,
     ) -> None:
         """
@@ -170,12 +171,25 @@ def fit(  # pylint: disable=W0221
                           the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                           the last batch will be smaller. (default: ``False``)
         :param scheduler: Learning rate scheduler to run at the start of every epoch.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                 class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                and providing it takes no effect.
         """
         import torch
         from torch.utils.data import TensorDataset, DataLoader
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -197,7 +211,7 @@ def fit(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
 
         # Start training
-        for epoch in trange(nb_epochs, disable=not self.verbose):
+        for epoch in trange(nb_epochs, disable=not display_pb):
             self.attack.norm = min(self.epsilon, (epoch + 1) * self.epsilon / self.warmup)
 
             for x_batch, y_batch in dataloader:

art/estimators/certification/randomized_smoothing/smooth_adv/tensorflow.py

@@ -149,7 +149,15 @@ def __init__(
         )
         self.attack = ProjectedGradientDescent(classifier, eps=self.epsilon, max_iter=1, verbose=False)
 
-    def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 10, **kwargs) -> None:
+    def fit(
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+        batch_size: int = 128,
+        nb_epochs: int = 10,
+        verbose: Optional[Union[bool, int]] = None,
+        **kwargs
+    ) -> None:
         """
         Fit the classifier on the training set `(x, y)`.
 
@@ -158,12 +166,25 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                   shape (nb_samples,).
         :param batch_size: Size of batches.
         :param nb_epochs: Number of epochs to use for training.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                 class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
                        "scheduler" which is an optional function that will be called at the end of every
                        epoch to adjust the learning rate.
         """
         import tensorflow as tf
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         if self._train_step is None:  # pragma: no cover
             if self._loss_object is None:  # pragma: no cover
                 raise TypeError(
@@ -200,7 +221,7 @@ def train_step(model, images, labels):
 
         train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)
 
-        for epoch in trange(nb_epochs, disable=not self.verbose):
+        for epoch in trange(nb_epochs, disable=not display_pb):
             self.attack.norm = min(self.epsilon, (epoch + 1) * self.epsilon / self.warmup)
 
             for x_batch, y_batch in train_ds:

art/estimators/certification/randomized_smoothing/smooth_mix/pytorch.py

@@ -172,6 +172,7 @@ def fit(  # pylint: disable=W0221
         training_mode: bool = True,
         drop_last: bool = False,
         scheduler: Optional["torch.optim.lr_scheduler._LRScheduler"] = None,
+        verbose: Optional[Union[bool, int]] = None,
         **kwargs,
     ) -> None:
         """
@@ -187,13 +188,26 @@ def fit(  # pylint: disable=W0221
                           the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                           the last batch will be smaller. (default: ``False``)
         :param scheduler: Learning rate scheduler to run at the start of every epoch.
+        :param verbose: (Optional) Display the progress bar, if not supplied will revert to the verbose level when
+                 class was initialised.
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                and providing it takes no effect.
         """
         import torch
         import torch.nn.functional as F
         from torch.utils.data import TensorDataset, DataLoader
 
+        if verbose is not None:
+            if isinstance(verbose, int):
+                if verbose == 0:
+                    display_pb = False
+                else:
+                    display_pb = True
+            else:
+                display_pb = verbose
+        else:
+            display_pb = self.verbose
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -215,7 +229,7 @@ def fit(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
 
         # Start training
-        for epoch in trange(nb_epochs, disable=not self.verbose):
+        for epoch in trange(nb_epochs, disable=not display_pb):
             warmup_v = min(1.0, (epoch + 1) / self.warmup)
 
             for x_batch, y_batch in dataloader: