Fix style checks

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

Author: Beat Buesser

art/estimators/certification/derandomized_smoothing/pytorch.py

@@ -147,82 +147,3 @@ def predict(
     def _fit_classifier(self, x: np.ndarray, y: np.ndarray, batch_size: int, nb_epochs: int, **kwargs) -> None:
         x = x.astype(ART_NUMPY_DTYPE)
         return PyTorchClassifier.fit(self, x, y, batch_size=batch_size, nb_epochs=nb_epochs, **kwargs)
-
-    def fit(  # pylint: disable=W0221
-        self,
-        x: np.ndarray,
-        y: np.ndarray,
-        batch_size: int = 128,
-        nb_epochs: int = 10,
-        training_mode: bool = True,
-        scheduler: Optional[Any] = None,
-        **kwargs,
-    ) -> None:
-        """
-        Fit the classifier on the training set `(x, y)`.
-
-        :param x: Training data.
-        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
-                  shape (nb_samples,).
-        :param batch_size: Size of batches.
-        :param nb_epochs: Number of epochs to use for training.
-        :param training_mode: `True` for model set to training mode and `'False` for model set to evaluation mode.
-        :param scheduler: Learning rate scheduler to run at the start of every epoch.
-        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
-               and providing it takes no effect.
-        """
-        import torch  # lgtm [py/repeated-import]
-
-        # Set model mode
-        self._model.train(mode=training_mode)
-
-        if self._optimizer is None:  # pragma: no cover
-            raise ValueError("An optimizer is needed to train the model, but none for provided.")
-
-        y = check_and_transform_label_format(y, nb_classes=self.nb_classes)
-
-        # Apply preprocessing
-        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=True)
-
-        # Check label shape
-        y_preprocessed = self.reduce_labels(y_preprocessed)
-
-        num_batch = int(np.ceil(len(x_preprocessed) / float(batch_size)))
-        ind = np.arange(len(x_preprocessed))
-
-        # Start training
-        for _ in tqdm(range(nb_epochs)):
-            # Shuffle the examples
-            random.shuffle(ind)
-
-            # Train for one epoch
-            for m in range(num_batch):
-                i_batch = np.copy(x_preprocessed[ind[m * batch_size : (m + 1) * batch_size]])
-                i_batch = self.ablator.forward(i_batch)
-
-                i_batch = torch.from_numpy(i_batch).to(self._device)
-                o_batch = torch.from_numpy(y_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]).to(self._device)
-
-                # Zero the parameter gradients
-                self._optimizer.zero_grad()
-
-                # Perform prediction
-                model_outputs = self._model(i_batch)
-
-                # Form the loss function
-                loss = self._loss(model_outputs[-1], o_batch)  # lgtm [py/call-to-non-callable]
-
-                # Do training
-                if self._use_amp:  # pragma: no cover
-                    from apex import amp  # pylint: disable=E0611
-
-                    with amp.scale_loss(loss, self._optimizer) as scaled_loss:
-                        scaled_loss.backward()
-
-                else:
-                    loss.backward()
-
-                self._optimizer.step()
-
-            if scheduler is not None:
-                scheduler.step()

art/estimators/classification/pytorch.py

@@ -363,6 +363,7 @@ def fit(  # pylint: disable=W0221
         nb_epochs: int = 10,
         training_mode: bool = True,
         drop_last: bool = False,
+        scheduler: Optional[Any] = None,
         **kwargs,
     ) -> None:
         """
@@ -377,6 +378,7 @@ def fit(  # pylint: disable=W0221
         :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                           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 kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                        and providing it takes no effect.
         """
@@ -422,13 +424,13 @@ def fit(  # pylint: disable=W0221
                 # Perform prediction
                 try:
                     model_outputs = self._model(i_batch)
-                except ValueError as e:
-                    if "Expected more than 1 value per channel when training" in str(e):
+                except ValueError as err:
+                    if "Expected more than 1 value per channel when training" in str(err):
                         logger.exception(
                             "Try dropping the last incomplete batch by setting drop_last=True in "
                             "method PyTorchClassifier.fit."
                         )
-                    raise e
+                    raise err
 
                 # Form the loss function
                 loss = self._loss(model_outputs[-1], o_batch)  # lgtm [py/call-to-non-callable]
@@ -445,6 +447,9 @@ def fit(  # pylint: disable=W0221
 
                 self._optimizer.step()
 
+            if scheduler is not None:
+                scheduler.step()
+
     def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwargs) -> None:
         """
         Fit the classifier using the generator that yields batches as specified.