[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

src/lightning/pytorch/callbacks/differential_privacy.py

@@ -52,7 +52,7 @@ def _set_layer_param(
     param: torch.Tensor,
     layer_name: str,
 ) -> None:
-    """Helper"""
+    """Helper."""
     layer = getattr(dpgrucell, layer_name)
     if "weight" in name:
         layer.weight = torch.nn.Parameter(deepcopy(param))
@@ -94,7 +94,7 @@ def params(
 
 
 class DPOptimizer(OpacusDPOptimizer):
-    """Brainiac-2's DP-Optimizer"""
+    """Brainiac-2's DP-Optimizer."""
 
     def __init__(
         self,
@@ -108,21 +108,20 @@ def __init__(
 
     @property
     def params(self) -> list[torch.nn.Parameter]:
-        """
-        Returns a flat list of ``nn.Parameter`` managed by the optimizer
-        """
+        """Returns a flat list of ``nn.Parameter`` managed by the optimizer."""
         return params(self, self.param_group_names)
 
 
 class DifferentialPrivacy(pl.callbacks.EarlyStopping):
-    """Enables differential privacy using Opacus.
-    Converts optimizers to instances of the :class:`~opacus.optimizers.DPOptimizer` class.
-    This callback inherits from `EarlyStopping`, thus it is also able to stop the
-    training when enough privacy budget has been spent.
-    Please beware that Opacus does not support multi-optimizer training.
+    """Enables differential privacy using Opacus. Converts optimizers to instances of the
+    :class:`~opacus.optimizers.DPOptimizer` class. This callback inherits from `EarlyStopping`, thus it is also able to
+    stop the training when enough privacy budget has been spent. Please beware that Opacus does not support multi-
+    optimizer training.
+
     For more info, check the following links:
     * https://opacus.ai/tutorials/
     * https://blog.openmined.org/differentially-private-deep-learning-using-opacus-in-20-lines-of-code/
+
     """
 
     def __init__(
@@ -140,6 +139,7 @@ def __init__(
         **gsm_kwargs: ty.Any,
     ) -> None:
         """Enables differential privacy using Opacus.
+
         Converts optimizers to instances of the :class:`~opacus.optimizers.DPOptimizer` class.
         This callback inherits from `EarlyStopping`,
         thus it is also able to stop the training when enough privacy budget has been spent.
@@ -177,6 +177,7 @@ def __init__(
                 Whether to make the dataloader private. Defaults to False.
             **gsm_kwargs:
                 Input arguments for the :class:`~opacus.GradSampleModule` class.
+
         """
         # inputs
         self.budget = budget
@@ -227,7 +228,11 @@ def on_train_epoch_start(
         trainer: pl.Trainer,
         pl_module: pl.LightningModule,
     ) -> None:
-        """Called when the training epoch begins. Use this to make optimizers private."""
+        """Called when the training epoch begins.
+
+        Use this to make optimizers private.
+
+        """
         # idx
         if self.idx is None:
             self.idx = range(len(trainer.optimizers))
@@ -289,7 +294,11 @@ def on_train_batch_end(  # pylint: disable=unused-argument # type: ignore
         batch_idx: int,
         *args: ty.Any,
     ) -> None:
-        """Called after the batched has been digested. Use this to understand whether to stop or not."""
+        """Called after the batched has been digested.
+
+        Use this to understand whether to stop or not.
+
+        """
         self._log_and_stop_criterion(trainer, pl_module)
 
     def on_train_epoch_end(

tests/tests_pytorch/callbacks/test_dp.py

@@ -68,6 +68,7 @@ def test_privacy_callback() -> None:
     * the privacy budget has been spent (`epsilon > 0`);
     * spent budget is  greater than max privacy budget;
     * traininng did not stop because `max_steps` has been reached, but because the total budget has been spent.
+
     """
     # choose dataset
     datamodule = MockDataModule()
@@ -92,12 +93,12 @@ def test_privacy_callback() -> None:
     epsilon, best_alpha = dp_cb.get_privacy_spent()
     print(f"Total spent budget {epsilon} with alpha: {best_alpha}")
     assert epsilon > 0, f"No privacy budget has been spent: {epsilon}"
-    assert (
-        epsilon >= dp_cb.budget
-    ), f"Spent budget is not greater than max privacy budget: epsilon = {epsilon} and budget = {dp_cb.budget}"
-    assert (
-        trainer.global_step < max_steps
-    ), "Traininng stopped because max_steps has been reached, not because the total budget has been spent."
+    assert epsilon >= dp_cb.budget, (
+        f"Spent budget is not greater than max privacy budget: epsilon = {epsilon} and budget = {dp_cb.budget}"
+    )
+    assert trainer.global_step < max_steps, (
+        "Traininng stopped because max_steps has been reached, not because the total budget has been spent."
+    )
 
 
 if __name__ == "__main__":