pep8

Co-authored-by: @awaelchi

Author: justusschock

pytorch_lightning/accelerators/accelerator.py

@@ -229,7 +229,14 @@ def backward(
 
         return output
 
-    def optimizer_step(self, optimizer: torch.optim.Optimizer, current_epoch: int, batch_idx: int, opt_idx: int, lambda_closure: Callable):
+    def optimizer_step(
+        self,
+        optimizer: torch.optim.Optimizer,
+        current_epoch: int,
+        batch_idx: int,
+        opt_idx: int,
+        lambda_closure: Callable,
+    ):
         """performs the actual optimizer step.
 
         Args:
@@ -265,16 +272,15 @@ def optimizer_step(self, optimizer: torch.optim.Optimizer, current_epoch: int, b
         self.training_type_plugin.post_optimizer_step(optimizer, opt_idx)
         return res
 
-    def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: torch.optim.Optimizer, opt_idx: int) -> None:
-        """Zeros all model parameter's gradients
-        """
+    def optimizer_zero_grad(
+        self, current_epoch: int, batch_idx: int, optimizer: torch.optim.Optimizer, opt_idx: int
+    ) -> None:
+        """Zeros all model parameter's gradients"""
         model_ref = self.lightning_module
         model_ref.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx)
 
     def clip_gradients(self, optimizer: torch.optim.Optimizer, clip_val: Union[int, float]) -> None:
-        """clips all the optimizer parameters to the given value
-
-        """
+        """clips all the optimizer parameters to the given value"""
 
         self.precision_plugin.clip_gradients(optimizer, clip_val)
 
@@ -287,11 +293,10 @@ def on_train_epoch_end(self, outputs) -> None:
         pass
 
     def on_train_end(self) -> None:
-        """Hook to do something at the end of the training
-        """
+        """Hook to do something at the end of the training"""
         pass
 
-    def setup_optimizers(self, trainer: 'Trainer', model: LightningModule):
+    def setup_optimizers(self, trainer: "Trainer", model: LightningModule):
         """creates optimizers and schedulers
 
         Args:
@@ -306,25 +311,21 @@ def setup_optimizers(self, trainer: 'Trainer', model: LightningModule):
         self.optimizer_frequencies = optimizer_frequencies
 
     def connect_training_type_plugin(self, plugin: TrainingTypePlugin, model: LightningModule) -> None:
-        """Attaches the training type plugin to the accelerator. 
+        """Attaches the training type plugin to the accelerator.
         Also transfers ownership of the model to this plugin
 
         """
         plugin.connect(model)
 
     def connect_precision_plugin(self, plugin: PrecisionPlugin):
-        """Attaches the precision plugin to the accelerator
-
-        """
+        """Attaches the precision plugin to the accelerator"""
         model, optimizers, schedulers = plugin.connect(self.model, self.optimizers, self.lr_schedulers)
         self.model = model
         self.optimizers = optimizers
         self.schedulers = schedulers
 
     def to_device(self, batch: Any) -> Any:
-        """Pushes the batch to the root device
-
-        """
+        """Pushes the batch to the root device"""
         return self.batch_to_device(batch, self.root_device)
 
     @property

pytorch_lightning/plugins/base_plugin.py

@@ -4,8 +4,7 @@
 
 
 class Plugin(object):
-    """Basic Plugin class to derive precision and training type plugins from.
-    """
+    """Basic Plugin class to derive precision and training type plugins from."""
 
     def connect(self, model: torch.nn.Module, *args, **kwargs):
         """Connects the plugin with the accelerator (and thereby with trainer and model).
@@ -14,39 +13,32 @@ def connect(self, model: torch.nn.Module, *args, **kwargs):
         pass
 
     def pre_optimizer_step(self, optimizer: torch.optim.Optimizer, optimizer_idx: int):
-        """Hook to do something before each optimizer step.
-        """
+        """Hook to do something before each optimizer step."""
         pass
 
     def post_optimizer_step(self, optimizer: torch.optim.Optimizer, optimizer_idx: int):
-        """Hook to do something after each optimizer step.
-        """
+        """Hook to do something after each optimizer step."""
         pass
 
     def pre_training(self):
-        """Hook to do something before the training starts.
-        """
+        """Hook to do something before the training starts."""
         pass
 
     def post_training(self):
-        """Hook to do something after the training finishes.
-        """
+        """Hook to do something after the training finishes."""
         pass
 
     @contextlib.contextmanager
     def train_step_context(self):
-        """A contextmanager for the trainstep
-        """
+        """A contextmanager for the trainstep"""
         yield
 
     @contextlib.contextmanager
     def val_step_context(self):
-        """A contextmanager for the validation step
-        """
+        """A contextmanager for the validation step"""
         yield
 
     @contextlib.contextmanager
     def test_step_context(self):
-        """A contextmanager for the teststep
-        """
+        """A contextmanager for the teststep"""
         yield

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -4,7 +4,7 @@
 import torch
 from torch.optim import Optimizer
 
-from pytorch_lightning.plugins .base_plugin import Plugin
+from pytorch_lightning.plugins.base_plugin import Plugin
 from pytorch_lightning.core import LightningModule
 
 
@@ -13,7 +13,7 @@ class PrecisionPlugin(Plugin):
     precision = 32
 
     def master_params(self, optimizer: torch.optim.Optimizer) -> Generator[torch.Tensor, None, None]:
-        """The master params of the model. Returns the plain model params here. 
+        """The master params of the model. Returns the plain model params here.
         Maybe different in other precision plugins.
 
         """
@@ -22,9 +22,7 @@ def master_params(self, optimizer: torch.optim.Optimizer) -> Generator[torch.Ten
                 yield p
 
     def connect(self, model: torch.nn.Module, optimizers, lr_schedulers):
-        """Connects this plugin to the accelerator and the training process
-
-        """
+        """Connects this plugin to the accelerator and the training process"""
         return model, optimizers, lr_schedulers
 
     def backward(
@@ -61,9 +59,7 @@ def backward(
         return closure_loss
 
     def clip_gradients(self, optimizer: Optimizer, clip_val: Union[int, float], norm_type: float = float(2.0)):
-        """Clips the gradients to a specific value
-
-        """
+        """Clips the gradients to a specific value"""
         # TODO: separate TPU case from here
         if clip_val is None:
             return

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -10,9 +10,8 @@
 
 
 class TrainingTypePlugin(Plugin, ABC):
-    """A Plugin to change the behaviour of the training, validation and test-loop.
+    """A Plugin to change the behaviour of the training, validation and test-loop."""
 
-    """
     def __init__(self):
         self._model = None
         self._results = None
@@ -21,46 +20,39 @@ def __init__(self):
     @property
     @abstractmethod
     def on_gpu(self) -> bool:
-        """Returns whether the current process is done on GPU
-        """
+        """Returns whether the current process is done on GPU"""
         raise NotImplementedError
 
     @property
     @abstractmethod
     def root_device(self) -> torch.device:
-        """Returns the root device
-        """
+        """Returns the root device"""
         raise NotImplementedError
 
     @abstractmethod
     def model_to_device(self):
-        """Moves the model to the correct device
-        """
+        """Moves the model to the correct device"""
         raise NotImplementedError
 
     @property
     @abstractmethod
     def is_global_zero(self) -> bool:
-        """Whether the current process is the rank zero process not only on the local node, but for all nodes.
-        """
+        """Whether the current process is the rank zero process not only on the local node, but for all nodes."""
         raise NotImplementedError
 
     @abstractmethod
     def reduce(self, output, *args, **kwargs):
-        """Reduces the given output (e.g. across GPUs/Processes)
-        """
+        """Reduces the given output (e.g. across GPUs/Processes)"""
         raise NotImplementedError
 
     @abstractmethod
     def barrier(self, name: Optional[str] = None):
-        """Forces all possibly joined processes to wait for each other
-        """
+        """Forces all possibly joined processes to wait for each other"""
         raise NotImplementedError
 
     @abstractmethod
     def broadcast(self, obj: object, src: int = 0) -> object:
-        """Broadcasts an object to all processes
-        """
+        """Broadcasts an object to all processes"""
         raise NotImplementedError
 
     # TODO method this is currently unused. Check after complete refactors are pushed
@@ -72,18 +64,15 @@ def set_nvidia_flags(self, is_slurm_managing_tasks, device_ids):
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
         all_gpu_ids = ",".join([str(x) for x in range(torch.cuda.device_count())])
         devices = os.environ.get("CUDA_VISIBLE_DEVICES", all_gpu_ids)
-        log.info(f'LOCAL_RANK: {self.trainer.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]')
+        log.info(f"LOCAL_RANK: {self.trainer.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]")
 
     def reduce_early_stopping_decision(self, should_stop: bool) -> bool:
-        """Reduce the early stopping decision across all possibly spawned processes
-        """
+        """Reduce the early stopping decision across all possibly spawned processes"""
         return should_stop
 
     @property
     def model(self) -> torch.nn.Module:
-        """Returns the potentially wrapped LightningModule
-
-        """
+        """Returns the potentially wrapped LightningModule"""
         return self._model
 
     @model.setter
@@ -92,9 +81,7 @@ def model(self, new_model: torch.nn.Module):
 
     @property
     def lightning_module(self) -> LightningModule:
-        """Returns the pure LightningModule without potential wrappers
-
-        """
+        """Returns the pure LightningModule without potential wrappers"""
         return self._model
 
     @property
@@ -110,10 +97,10 @@ def results(self):
     def rpc_enabled(self) -> bool:
         return False
 
-    def start_training(self, trainer: 'Trainer') -> None:
+    def start_training(self, trainer: "Trainer") -> None:
         # double dispatch to initiate the training loop
         self._results = trainer.train()
 
-    def start_testing(self, trainer: 'Trainer') -> None:
+    def start_testing(self, trainer: "Trainer") -> None:
         # double dispatch to initiate the test loop
         self._results = trainer.run_test()