Merge branch 'master' into refactor/mc-save-on-train-epoch-end

Author: carmocca

.github/workflows/code-checks.yml

@@ -17,6 +17,6 @@ jobs:
           python-version: 3.8
       - name: Install mypy
         run: |
-          pip install mypy==0.790
+          grep mypy requirements/test.txt | xargs -0 pip install
           pip list
       - run: mypy

.pre-commit-config.yaml

@@ -15,6 +15,12 @@
 default_language_version:
   python: python3.8
 
+ci:
+  autofix_prs: true
+  autoupdate_commit_msg: '[pre-commit.ci] pre-commit suggestions'
+  autoupdate_schedule: quarterly
+  # submodules: true
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.0.1
@@ -40,7 +46,7 @@ repos:
       - id: detect-private-key
 
   - repo: https://github.com/PyCQA/isort
-    rev: 5.9.1
+    rev: 5.9.2
     hooks:
       - id: isort
         name: Format imports

CHANGELOG.md

@@ -9,6 +9,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Add `extract_batch_size` utility and corresponding tests to extract batch dimension from multiple batch types. ([#8357](https://github.com/PyTorchLightning/pytorch-lightning/pull/8357/))
+
 - Add support for named parameter groups in `LearningRateMonitor` ([#7987](https://github.com/PyTorchLightning/pytorch-lightning/pull/7987))
 
 
@@ -116,7 +118,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support `LightningModule.save_hyperparameters` when `LightningModule` is a dataclass ([#7992](https://github.com/PyTorchLightning/pytorch-lightning/pull/7992))
 
 
-- Add support for overriding `optimizer_zero_grad` and `optimizer_step` when using accumulate_grad_batches ([#7980](https://github.com/PyTorchLightning/pytorch-lightning/pull/7980))
+- Added support for overriding `optimizer_zero_grad` and `optimizer_step` when using accumulate_grad_batches ([#7980](https://github.com/PyTorchLightning/pytorch-lightning/pull/7980))
+
+
+- Added `logger` boolean flag to `save_hyperparameters` ([#7960](https://github.com/PyTorchLightning/pytorch-lightning/pull/7960))
 
 
 - Add support for calling scripts using the module syntax (`python -m package.script`) ([#8073](https://github.com/PyTorchLightning/pytorch-lightning/pull/8073))

pyproject.toml

@@ -4,6 +4,7 @@ requires = [
     "wheel",
 ]
 
+
 [tool.isort]
 known_first_party = [
     "benchmarks",
@@ -16,3 +17,51 @@ profile = "black"
 line_length = 120
 force_sort_within_sections = "False"
 order_by_type = "False"
+
+
+[tool.mypy]
+files = ["pytorch_lightning", "pl_examples", "benchmarks"]
+disallow_untyped_defs = "True"
+ignore_missing_imports = "True"
+show_error_codes = "True"
+warn_redundant_casts = "True"
+warn_unused_configs = "True"
+warn_unused_ignores = "True"
+allow_redefinition = "True"
+# disable this rule as the Trainer attributes are defined in the connectors, not in its __init__
+disable_error_code = "attr-defined"
+
+
+# TODO: Fix typing for these modules
+[[tool.mypy.overrides]]
+module = [
+    "pytorch_lightning.callbacks.*",
+    "pytorch_lightning.core.*",
+    "pytorch_lightning.loggers.*",
+    "pytorch_lightning.loops.*",
+    "pytorch_lightning.metrics.*",
+    "pytorch_lightning.overrides.*",
+    "pytorch_lightning.plugins.environments.*",
+    "pytorch_lightning.plugins.training_type.*",
+    "pytorch_lightning.profiler.*",
+    "pytorch_lightning.trainer.*",
+    "pytorch_lightning.distributed.*",
+    "pytorch_lightning.tuner.*",
+    "pytorch_lightning.utilities.*",
+    "pl_examples.*",
+    "benchmarks.*",
+    "tests.helpers.*"
+]
+ignore_errors = "True"
+
+[[tool.mypy.overrides]]
+module = [
+    "pytorch_lightning.callbacks.pruning",
+    "pytorch_lightning.trainer.evaluation_loop",
+    "pytorch_lightning.trainer.connectors.logger_connector",
+    "pytorch_lightning.utilities.cli",
+    "pytorch_lightning.utilities.device_dtype_mixin",
+    "pytorch_lightning.utilities.device_parser",
+    "pytorch_lightning.utilities.parsing",
+]
+ignore_errors = "False"

pytorch_lightning/accelerators/accelerator.py

@@ -174,6 +174,8 @@ def batch_to_device(
             dataloader_idx: The index of the dataloader to which the batch belongs.
         """
         model = self.lightning_module
+        device = device or self.root_device
+
         if model is not None and not isinstance(self.training_type_plugin, DataParallelPlugin):
             # no need to transfer batch to device in DP mode
             return model._apply_batch_transfer_handler(batch, device, dataloader_idx)

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -46,7 +46,7 @@ def __init__(self, trainer: 'pl.Trainer', log_gpu_memory: Optional[str] = None)
 
     def on_trainer_init(
         self,
-        logger: LightningLoggerBase,
+        logger: Union[bool, LightningLoggerBase, Iterable[LightningLoggerBase]],
         flush_logs_every_n_steps: int,
         log_every_n_steps: int,
         move_metrics_to_cpu: bool,
@@ -66,7 +66,7 @@ def should_update_logs(self) -> bool:
         should_log_every_n_steps = (self.trainer.global_step + 1) % self.trainer.log_every_n_steps == 0
         return should_log_every_n_steps or self.trainer.should_stop
 
-    def configure_logger(self, logger: Union[bool, Iterable, LightningLoggerBase]) -> None:
+    def configure_logger(self, logger: Union[bool, LightningLoggerBase, Iterable[LightningLoggerBase]]) -> None:
         if logger is True:
             version = os.environ.get('PL_EXP_VERSION', self.trainer.slurm_job_id)
 

pytorch_lightning/trainer/connectors/logger_connector/result.py

@@ -14,13 +14,14 @@
 from collections.abc import Generator
 from dataclasses import asdict, dataclass, replace
 from functools import partial, wraps
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
 
 import torch
 from torchmetrics import Metric
 
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection, apply_to_collections
+from pytorch_lightning.utilities.data import extract_batch_size
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 from pytorch_lightning.utilities.distributed import distributed_available
 from pytorch_lightning.utilities.enums import LightningEnum
@@ -589,31 +590,10 @@ def fn(item: ResultMetric) -> None:
 
     def extract_batch_size(self, batch: Any) -> None:
         try:
-            self.batch_size = self._extract_batch_size(batch)
+            self.batch_size = extract_batch_size(batch)
         except RecursionError:
             self.batch_size = 1
 
-    def _extract_batch_size(self, batch: Any) -> int:
-        """
-        Recursively unpack a batch to find a torch.Tensor.
-
-        Returns:
-            ``len(tensor)`` when found, or ``1`` when it hits an empty or non iterable.
-        """
-        if isinstance(batch, torch.Tensor):
-            size = batch.size(0)
-        elif isinstance(batch, str):
-            return len(batch)
-        elif isinstance(batch, dict):
-            sample = next(iter(batch.values()), 1)
-            size = self._extract_batch_size(sample)
-        elif isinstance(batch, Iterable):
-            sample = next(iter(batch), 1)
-            size = self._extract_batch_size(sample)
-        else:
-            size = 1
-        return size
-
     def to(self, *args, **kwargs) -> 'ResultCollection':
         """Move all data to the given device."""
 

pytorch_lightning/trainer/trainer.py

@@ -907,20 +907,30 @@ def _pre_dispatch(self):
 
     def _log_hyperparams(self):
         # log hyper-parameters
+        hparams_initial = None
+
         if self.logger is not None:
             # save exp to get started (this is where the first experiment logs are written)
-            datamodule_hparams = self.datamodule.hparams_initial if self.datamodule is not None else {}
-            lightning_hparams = self.lightning_module.hparams_initial
-            colliding_keys = lightning_hparams.keys() & datamodule_hparams.keys()
-            if colliding_keys:
-                raise MisconfigurationException(
-                    f"Error while merging hparams: the keys {colliding_keys} are present "
-                    "in both the LightningModule's and LightningDataModule's hparams."
-                )
+            datamodule_log_hyperparams = self.datamodule._log_hyperparams if self.datamodule is not None else False
 
-            hparams_initial = {**lightning_hparams, **datamodule_hparams}
+            if self.lightning_module._log_hyperparams and datamodule_log_hyperparams:
+                datamodule_hparams = self.datamodule.hparams_initial
+                lightning_hparams = self.lightning_module.hparams_initial
 
-            self.logger.log_hyperparams(hparams_initial)
+                colliding_keys = lightning_hparams.keys() & datamodule_hparams.keys()
+                if colliding_keys:
+                    raise MisconfigurationException(
+                        f"Error while merging hparams: the keys {colliding_keys} are present "
+                        "in both the LightningModule's and LightningDataModule's hparams."
+                    )
+                hparams_initial = {**lightning_hparams, **datamodule_hparams}
+            elif self.lightning_module._log_hyperparams:
+                hparams_initial = self.lightning_module.hparams_initial
+            elif datamodule_log_hyperparams:
+                hparams_initial = self.datamodule.hparams_initial
+
+            if hparams_initial is not None:
+                self.logger.log_hyperparams(hparams_initial)
             self.logger.log_graph(self.lightning_module)
             self.logger.save()
 

pytorch_lightning/utilities/data.py

@@ -12,12 +12,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from typing import Any, Iterable, Mapping, Union
 
+import torch
 from torch.utils.data import DataLoader, IterableDataset
 
 from pytorch_lightning.utilities import rank_zero_warn
 
+BType = Union[torch.Tensor, str, Mapping[Any, 'BType'], Iterable['BType']]
+
+
+def extract_batch_size(batch: BType) -> int:
+    """
+    Recursively unpack a batch to find a torch.Tensor.
+
+    Returns:
+        ``len(tensor)`` when found, or ``1`` when it hits an empty or non iterable.
+    """
+    if isinstance(batch, torch.Tensor):
+        return batch.size(0)
+    if isinstance(batch, str):
+        return len(batch)
+    if isinstance(batch, dict):
+        sample = next(iter(batch.values()), 1)
+        return extract_batch_size(sample)
+    if isinstance(batch, Iterable):
+        sample = next(iter(batch), 1)
+        return extract_batch_size(sample)
+
+    return 1
+
 
 def has_iterable_dataset(dataloader: DataLoader):
     return hasattr(dataloader, 'dataset') and isinstance(dataloader.dataset, IterableDataset)

pytorch_lightning/utilities/device_dtype_mixin.py

@@ -121,10 +121,9 @@ def cuda(self, device: Optional[Union[torch.device, int]] = None) -> 'DeviceDtyp
         Returns:
             Module: self
         """
-        property_device = (
-            device if isinstance(device, torch.device) else torch.device('cuda', index=device)  # type: ignore
-        )  # mypy expects `device` for `index` to be int, while `Optional[int]` is okay => ignore typing for now
-        self.__update_properties(device=property_device)
+        if device is None or isinstance(device, int):
+            device = torch.device('cuda', index=device)
+        self.__update_properties(device=device)
         return super().cuda(device=device)
 
     def cpu(self) -> 'DeviceDtypeModuleMixin':