-- [Napoleon formatting with Google style](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
-- [ReStructured Text (reST)](https://docs.pylonsproject.org/projects/docs-style-guide/)
-- [Paragraph-level markup](https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html#paragraphs)
++ Lightning.ai • + PyTorch Lightning • + Fabric • + Lightning Apps • + Docs • + Community • + Contribute • +
+ + + +[](https://pypi.org/project/pytorch-lightning/) +[](https://badge.fury.io/py/pytorch-lightning) +[](https://pepy.tech/project/pytorch-lightning) +[](https://anaconda.org/conda-forge/lightning) +[](https://hub.docker.com/r/pytorchlightning/pytorch_lightning) +[](https://codecov.io/gh/Lightning-AI/lightning) + +[](https://discord.gg/VptPCZkGNa) + +[](https://github.com/Lightning-AI/lightning/blob/master/LICENSE) + + + +
+
+ | What to change | +Resulting Fabric Code (copy me!) | +
|---|---|
| + + +```diff ++ import lightning as L + import torch; import torchvision as tv + + dataset = tv.datasets.CIFAR10("data", download=True, + train=True, + transform=tv.transforms.ToTensor()) + ++ fabric = L.Fabric() ++ fabric.launch() + + model = tv.models.resnet18() + optimizer = torch.optim.SGD(model.parameters(), lr=0.001) +- device = "cuda" if torch.cuda.is_available() else "cpu" +- model.to(device) ++ model, optimizer = fabric.setup(model, optimizer) + + dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) ++ dataloader = fabric.setup_dataloaders(dataloader) + + model.train() + num_epochs = 10 + for epoch in range(num_epochs): + for batch in dataloader: + inputs, labels = batch +- inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + outputs = model(inputs) + loss = torch.nn.functional.cross_entropy(outputs, labels) +- loss.backward() ++ fabric.backward(loss) + optimizer.step() + print(loss.data) +``` + + + | + + +```Python +import lightning as L +import torch; import torchvision as tv + +dataset = tv.datasets.CIFAR10("data", download=True, + train=True, + transform=tv.transforms.ToTensor()) + +fabric = L.Fabric() +fabric.launch() + +model = tv.models.resnet18() +optimizer = torch.optim.SGD(model.parameters(), lr=0.001) +model, optimizer = fabric.setup(model, optimizer) + +dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) +dataloader = fabric.setup_dataloaders(dataloader) + +model.train() +num_epochs = 10 +for epoch in range(num_epochs): + for batch in dataloader: + inputs, labels = batch + optimizer.zero_grad() + outputs = model(inputs) + loss = torch.nn.functional.cross_entropy(outputs, labels) + fabric.backward(loss) + optimizer.step() + print(loss.data) +``` + + + | +
+
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-r"""
-Base class used to build new callbacks.
-
-"""
-
-from typing import Any, Dict, List, Optional, Type
-
-import torch
-from torch.optim import Optimizer
-
-import pytorch_lightning as pl
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-
-
-class Callback:
- r"""
- Abstract base class used to build new callbacks.
-
- Subclass this class and override any of the relevant hooks
- """
-
- @property
- def state_key(self) -> str:
- """Identifier for the state of the callback.
-
- Used to store and retrieve a callback's state from the checkpoint dictionary by
- ``checkpoint["callbacks"][state_key]``. Implementations of a callback need to provide a unique state key if 1)
- the callback has state and 2) it is desired to maintain the state of multiple instances of that callback.
- """
- return self.__class__.__qualname__
-
- @property
- def _legacy_state_key(self) -> Type["Callback"]:
- """State key for checkpoints saved prior to version 1.5.0."""
- return type(self)
-
- def _generate_state_key(self, **kwargs: Any) -> str:
- """Formats a set of key-value pairs into a state key string with the callback class name prefixed. Useful
- for defining a :attr:`state_key`.
-
- Args:
- **kwargs: A set of key-value pairs. Must be serializable to :class:`str`.
- """
- return f"{self.__class__.__qualname__}{repr(kwargs)}"
-
-[docs] def on_configure_sharded_model(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use `setup()` instead.
-
- Called before configure sharded model.
- """
-
- def on_before_accelerator_backend_setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use ``setup()`` instead.
-
- Called before accelerator is being setup.
- """
-
-[docs] def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: Optional[str] = None) -> None:
- """Called when fit, validate, test, predict, or tune begins."""
-
-[docs] def teardown(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: Optional[str] = None) -> None:
- """Called when fit, validate, test, predict, or tune ends."""
-
-[docs] def on_init_start(self, trainer: "pl.Trainer") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8.
-
- Called when the trainer initialization begins, model has not yet been set.
- """
-
-[docs] def on_init_end(self, trainer: "pl.Trainer") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8.
-
- Called when the trainer initialization ends, model has not yet been set.
- """
-
-[docs] def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when fit begins."""
-
-[docs] def on_fit_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when fit ends."""
-
-[docs] def on_sanity_check_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the validation sanity check starts."""
-
-[docs] def on_sanity_check_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the validation sanity check ends."""
-
-[docs] def on_train_batch_start(
- self,
- trainer: "pl.Trainer",
- pl_module: "pl.LightningModule",
- batch: Any,
- batch_idx: int,
- unused: int = 0,
- ) -> None:
- """Called when the train batch begins."""
-
-[docs] def on_train_batch_end(
- self,
- trainer: "pl.Trainer",
- pl_module: "pl.LightningModule",
- outputs: STEP_OUTPUT,
- batch: Any,
- batch_idx: int,
- unused: int = 0,
- ) -> None:
- """Called when the train batch ends."""
-
-[docs] def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the train epoch begins."""
-
-[docs] def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the train epoch ends.
-
- To access all batch outputs at the end of the epoch, either:
-
- 1. Implement `training_epoch_end` in the `LightningModule` and access outputs via the module OR
- 2. Cache data across train batch hooks inside the callback implementation to post-process in this hook.
- """
-
-[docs] def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the val epoch begins."""
-
-[docs] def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the val epoch ends."""
-
-[docs] def on_test_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the test epoch begins."""
-
-[docs] def on_test_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the test epoch ends."""
-
-[docs] def on_predict_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the predict epoch begins."""
-
-[docs] def on_predict_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: List[Any]) -> None:
- """Called when the predict epoch ends."""
-
- def on_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use
- ``on_<train/validation/test>_epoch_start`` instead.
-
- Called when either of train/val/test epoch begins.
- """
-
-[docs] def on_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use
- ``on_<train/validation/test>_epoch_end`` instead.
-
- Called when either of train/val/test epoch ends.
- """
-
- def on_batch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use
- ``on_train_batch_start`` instead.
-
- Called when the training batch begins.
- """
-
- def on_batch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use
- ``on_train_batch_end`` instead.
-
- Called when the training batch ends.
- """
-
-[docs] def on_validation_batch_start(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int
- ) -> None:
- """Called when the validation batch begins."""
-
-[docs] def on_validation_batch_end(
- self,
- trainer: "pl.Trainer",
- pl_module: "pl.LightningModule",
- outputs: Optional[STEP_OUTPUT],
- batch: Any,
- batch_idx: int,
- dataloader_idx: int,
- ) -> None:
- """Called when the validation batch ends."""
-
-[docs] def on_test_batch_start(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int
- ) -> None:
- """Called when the test batch begins."""
-
-[docs] def on_test_batch_end(
- self,
- trainer: "pl.Trainer",
- pl_module: "pl.LightningModule",
- outputs: Optional[STEP_OUTPUT],
- batch: Any,
- batch_idx: int,
- dataloader_idx: int,
- ) -> None:
- """Called when the test batch ends."""
-
-[docs] def on_predict_batch_start(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int
- ) -> None:
- """Called when the predict batch begins."""
-
-[docs] def on_predict_batch_end(
- self,
- trainer: "pl.Trainer",
- pl_module: "pl.LightningModule",
- outputs: Any,
- batch: Any,
- batch_idx: int,
- dataloader_idx: int,
- ) -> None:
- """Called when the predict batch ends."""
-
-[docs] def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the train begins."""
-
-[docs] def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the train ends."""
-
- def on_pretrain_routine_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
-
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use ``on_fit_start`` instead.
-
- Called when the pretrain routine begins.
- """
-
- def on_pretrain_routine_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.6
-
- This callback hook was deprecated in v1.6 and will be removed in v1.8. Use ``on_fit_start`` instead.
-
- Called when the pretrain routine ends.
- """
-
-[docs] def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the validation loop begins."""
-
-[docs] def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the validation loop ends."""
-
-[docs] def on_test_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the test begins."""
-
-[docs] def on_test_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the test ends."""
-
-[docs] def on_predict_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when the predict begins."""
-
-[docs] def on_predict_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called when predict ends."""
-
-[docs] def on_keyboard_interrupt(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- r"""
- .. deprecated:: v1.5
- This callback hook was deprecated in v1.5 in favor of `on_exception` and will be removed in v1.7.
-
- Called when any trainer execution is interrupted by KeyboardInterrupt.
- """
-
-[docs] def on_exception(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", exception: BaseException) -> None:
- """Called when any trainer execution is interrupted by an exception."""
-
-[docs] def state_dict(self) -> Dict[str, Any]:
- """Called when saving a checkpoint, implement to generate callback's ``state_dict``.
-
- Returns:
- A dictionary containing callback state.
- """
- return {}
-
-[docs] def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
- """Called when loading a checkpoint, implement to reload callback state given callback's ``state_dict``.
-
- Args:
- state_dict: the callback state returned by ``state_dict``.
- """
- pass
-
-[docs] def on_save_checkpoint(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
- ) -> Optional[dict]:
- r"""
- Called when saving a checkpoint to give you a chance to store anything else you might want to save.
-
- Args:
- trainer: the current :class:`~pytorch_lightning.trainer.Trainer` instance.
- pl_module: the current :class:`~pytorch_lightning.core.lightning.LightningModule` instance.
- checkpoint: the checkpoint dictionary that will be saved.
-
- Returns:
- None or the callback state. Support for returning callback state will be removed in v1.8.
-
- .. deprecated:: v1.6
- Returning a value from this method was deprecated in v1.6 and will be removed in v1.8.
- Implement ``Callback.state_dict`` instead to return state.
- In v1.8 ``Callback.on_save_checkpoint`` can only return None.
- """
-
-[docs] def on_load_checkpoint(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", callback_state: Dict[str, Any]
- ) -> None:
- r"""
- Called when loading a model checkpoint, use to reload state.
-
- Args:
- trainer: the current :class:`~pytorch_lightning.trainer.Trainer` instance.
- pl_module: the current :class:`~pytorch_lightning.core.lightning.LightningModule` instance.
- callback_state: the callback state returned by ``on_save_checkpoint``.
-
- Note:
- The ``on_load_checkpoint`` won't be called with an undefined state.
- If your ``on_load_checkpoint`` hook behavior doesn't rely on a state,
- you will still need to override ``on_save_checkpoint`` to return a ``dummy state``.
-
- .. deprecated:: v1.6
- This callback hook will change its signature and behavior in v1.8.
- If you wish to load the state of the callback, use ``Callback.load_state_dict`` instead.
- In v1.8 ``Callback.on_load_checkpoint(checkpoint)`` will receive the entire loaded
- checkpoint dictionary instead of only the callback state from the checkpoint.
- """
-
-[docs] def on_before_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", loss: torch.Tensor) -> None:
- """Called before ``loss.backward()``."""
-
-[docs] def on_after_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
- """Called after ``loss.backward()`` and before optimizers are stepped."""
-
-[docs] def on_before_optimizer_step(
- self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", optimizer: Optimizer, opt_idx: int
- ) -> None:
- """Called before ``optimizer.step()``."""
-
-[docs] def on_before_zero_grad(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", optimizer: Optimizer) -> None:
- """Called before ``optimizer.zero_grad()``."""
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""LightningDataModule for loading DataLoaders with ease."""
-from argparse import ArgumentParser, Namespace
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
-
-from torch.utils.data import DataLoader, Dataset, IterableDataset
-
-from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks
-from pytorch_lightning.core.mixins import HyperparametersMixin
-from pytorch_lightning.utilities import rank_zero_deprecation
-from pytorch_lightning.utilities.argparse import add_argparse_args, from_argparse_args, get_init_arguments_and_types
-
-
-class LightningDataModule(CheckpointHooks, DataHooks, HyperparametersMixin):
- """A DataModule standardizes the training, val, test splits, data preparation and transforms. The main
- advantage is consistent data splits, data preparation and transforms across models.
-
- Example::
-
- class MyDataModule(LightningDataModule):
- def __init__(self):
- super().__init__()
- def prepare_data(self):
- # download, split, etc...
- # only called on 1 GPU/TPU in distributed
- def setup(self, stage):
- # make assignments here (val/train/test split)
- # called on every process in DDP
- def train_dataloader(self):
- train_split = Dataset(...)
- return DataLoader(train_split)
- def val_dataloader(self):
- val_split = Dataset(...)
- return DataLoader(val_split)
- def test_dataloader(self):
- test_split = Dataset(...)
- return DataLoader(test_split)
- def teardown(self):
- # clean up after fit or test
- # called on every process in DDP
- """
-
- name: str = ...
-
- def __init__(self, train_transforms=None, val_transforms=None, test_transforms=None, dims=None):
- super().__init__()
- if train_transforms is not None:
- rank_zero_deprecation(
- "DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- if val_transforms is not None:
- rank_zero_deprecation(
- "DataModule property `val_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- if test_transforms is not None:
- rank_zero_deprecation(
- "DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- if dims is not None:
- rank_zero_deprecation("DataModule property `dims` was deprecated in v1.5 and will be removed in v1.7.")
- self._train_transforms = train_transforms
- self._val_transforms = val_transforms
- self._test_transforms = test_transforms
- self._dims = dims if dims is not None else ()
-
- # Pointer to the trainer object
- self.trainer = None
-
- @property
- def train_transforms(self):
- """Optional transforms (or collection of transforms) you can apply to train dataset.
-
- .. deprecated:: v1.5 Will be removed in v1.7.0.
- """
-
- rank_zero_deprecation(
- "DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- return self._train_transforms
-
- @train_transforms.setter
- def train_transforms(self, t):
- rank_zero_deprecation(
- "DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- self._train_transforms = t
-
- @property
- def val_transforms(self):
- """Optional transforms (or collection of transforms) you can apply to validation dataset.
-
- .. deprecated:: v1.5 Will be removed in v1.7.0.
- """
-
- rank_zero_deprecation(
- "DataModule property `val_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- return self._val_transforms
-
- @val_transforms.setter
- def val_transforms(self, t):
- rank_zero_deprecation(
- "DataModule property `val_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- self._val_transforms = t
-
- @property
- def test_transforms(self):
- """Optional transforms (or collection of transforms) you can apply to test dataset.
-
- .. deprecated:: v1.5 Will be removed in v1.7.0.
- """
-
- rank_zero_deprecation(
- "DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- return self._test_transforms
-
- @test_transforms.setter
- def test_transforms(self, t):
- rank_zero_deprecation(
- "DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7."
- )
- self._test_transforms = t
-
- @property
- def dims(self):
- """A tuple describing the shape of your data. Extra functionality exposed in ``size``.
-
- .. deprecated:: v1.5 Will be removed in v1.7.0.
- """
- rank_zero_deprecation("DataModule property `dims` was deprecated in v1.5 and will be removed in v1.7.")
- return self._dims
-
- @dims.setter
- def dims(self, d):
- rank_zero_deprecation("DataModule property `dims` was deprecated in v1.5 and will be removed in v1.7.")
- self._dims = d
-
- def size(self, dim=None) -> Union[Tuple, List[Tuple]]:
- """Return the dimension of each input either as a tuple or list of tuples. You can index this just as you
- would with a torch tensor.
-
- .. deprecated:: v1.5 Will be removed in v1.7.0.
- """
- rank_zero_deprecation("DataModule property `size` was deprecated in v1.5 and will be removed in v1.7.")
-
- if dim is not None:
- return self.dims[dim]
-
- return self.dims
-
- @classmethod
- def add_argparse_args(cls, parent_parser: ArgumentParser, **kwargs) -> ArgumentParser:
- """Extends existing argparse by default `LightningDataModule` attributes."""
- return add_argparse_args(cls, parent_parser, **kwargs)
-
- @classmethod
- def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):
- """Create an instance from CLI arguments.
-
- Args:
- args: The parser or namespace to take arguments from. Only known arguments will be
- parsed and passed to the :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
- **kwargs: Additional keyword arguments that may override ones in the parser or namespace.
- These must be valid DataModule arguments.
-
- Example::
-
- parser = ArgumentParser(add_help=False)
- parser = LightningDataModule.add_argparse_args(parser)
- module = LightningDataModule.from_argparse_args(args)
- """
- return from_argparse_args(cls, args, **kwargs)
-
- @classmethod
- def get_init_arguments_and_types(cls) -> List[Tuple[str, Tuple, Any]]:
- r"""Scans the DataModule signature and returns argument names, types and default values.
-
- Returns:
- List with tuples of 3 values:
- (argument name, set with argument types, argument default value).
- """
- return get_init_arguments_and_types(cls)
-
- @classmethod
- def from_datasets(
- cls,
- train_dataset: Optional[Union[Dataset, Sequence[Dataset], Mapping[str, Dataset]]] = None,
- val_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
- test_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
- batch_size: int = 1,
- num_workers: int = 0,
- ):
- r"""
- Create an instance from torch.utils.data.Dataset.
-
- Args:
- train_dataset: (optional) Dataset to be used for train_dataloader()
- val_dataset: (optional) Dataset or list of Dataset to be used for val_dataloader()
- test_dataset: (optional) Dataset or list of Dataset to be used for test_dataloader()
- batch_size: Batch size to use for each dataloader. Default is 1.
- num_workers: Number of subprocesses to use for data loading. 0 means that the
- data will be loaded in the main process. Number of CPUs available.
-
- """
-
- def dataloader(ds: Dataset, shuffle: bool = False) -> DataLoader:
- shuffle &= not isinstance(ds, IterableDataset)
- return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True)
-
- def train_dataloader():
- if isinstance(train_dataset, Mapping):
- return {key: dataloader(ds, shuffle=True) for key, ds in train_dataset.items()}
- if isinstance(train_dataset, Sequence):
- return [dataloader(ds, shuffle=True) for ds in train_dataset]
- return dataloader(train_dataset, shuffle=True)
-
- def val_dataloader():
- if isinstance(val_dataset, Sequence):
- return [dataloader(ds) for ds in val_dataset]
- return dataloader(val_dataset)
-
- def test_dataloader():
- if isinstance(test_dataset, Sequence):
- return [dataloader(ds) for ds in test_dataset]
- return dataloader(test_dataset)
-
- datamodule = cls()
- if train_dataset is not None:
- datamodule.train_dataloader = train_dataloader
- if val_dataset is not None:
- datamodule.val_dataloader = val_dataloader
- if test_dataset is not None:
- datamodule.test_dataloader = test_dataloader
- return datamodule
-
-[docs] def state_dict(self) -> Dict[str, Any]:
- """Called when saving a checkpoint, implement to generate and save datamodule state.
-
- Returns:
- A dictionary containing datamodule state.
- """
- return {}
-
-[docs] def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
- """Called when loading a checkpoint, implement to reload datamodule state given datamodule state_dict.
-
- Args:
- state_dict: the datamodule state returned by ``state_dict``.
- """
- pass
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""The LightningModule - an nn.Module with many additional features."""
-
-import collections
-import inspect
-import logging
-import numbers
-import os
-import tempfile
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Mapping, Optional, overload, Sequence, Tuple, Union
-
-import torch
-from torch import ScriptModule, Tensor
-from torch.nn import Module
-from torch.optim.optimizer import Optimizer
-from torchmetrics import Metric
-from typing_extensions import Literal
-
-import pytorch_lightning as pl
-from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.callbacks.progress import base as progress_base
-from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks, ModelHooks
-from pytorch_lightning.core.mixins import DeviceDtypeModuleMixin, HyperparametersMixin
-from pytorch_lightning.core.optimizer import LightningOptimizer
-from pytorch_lightning.core.saving import ModelIO
-from pytorch_lightning.loggers import LightningLoggerBase
-from pytorch_lightning.trainer.connectors.data_connector import _DataHookSelector
-from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import _FxValidator
-from pytorch_lightning.utilities import _IS_WINDOWS, _TORCH_GREATER_EQUAL_1_10, GradClipAlgorithmType
-from pytorch_lightning.utilities.apply_func import apply_to_collection, convert_to_tensors
-from pytorch_lightning.utilities.cloud_io import get_filesystem
-from pytorch_lightning.utilities.distributed import distributed_available, sync_ddp
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.memory import get_model_size_mb
-from pytorch_lightning.utilities.model_summary import ModelSummary, summarize
-from pytorch_lightning.utilities.parsing import collect_init_args
-from pytorch_lightning.utilities.rank_zero import rank_zero_debug, rank_zero_deprecation, rank_zero_warn
-from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
-from pytorch_lightning.utilities.types import _METRIC_COLLECTION, EPOCH_OUTPUT, LRSchedulerTypeUnion, STEP_OUTPUT
-from pytorch_lightning.utilities.warnings import WarningCache
-
-warning_cache = WarningCache()
-log = logging.getLogger(__name__)
-
-
-class LightningModule(
- DeviceDtypeModuleMixin,
- HyperparametersMixin,
- ModelIO,
- ModelHooks,
- DataHooks,
- CheckpointHooks,
- Module,
-):
- # Below is for property support of JIT
- # since none of these are important when using JIT, we are going to ignore them.
- __jit_unused_properties__ = (
- [
- "example_input_array",
- "on_gpu",
- "current_epoch",
- "global_step",
- "global_rank",
- "local_rank",
- "logger",
- "loggers",
- "model_size",
- "automatic_optimization",
- "truncated_bptt_steps",
- "use_amp",
- ]
- + DeviceDtypeModuleMixin.__jit_unused_properties__
- + HyperparametersMixin.__jit_unused_properties__
- )
-
- def __init__(self, *args: Any, **kwargs: Any) -> None:
- super().__init__(*args, **kwargs)
-
- # see (https://github.com/pytorch/pytorch/blob/3e6bb5233f9ca2c5aa55d9cda22a7ee85439aa6e/
- # torch/nn/modules/module.py#L227)
- torch._C._log_api_usage_once(f"lightning.module.{self.__class__.__name__}")
-
- # pointer to the trainer object
- self.trainer: Optional["pl.Trainer"] = None
-
- self._use_amp: bool = False
-
- # the precision used
- self.precision: int = 32
-
- # optionally can be set by user
- self._example_input_array = None
- self._current_fx_name: Optional[str] = None
- self._automatic_optimization: bool = True
- self._truncated_bptt_steps: int = 0
- self._param_requires_grad_state = {}
- self._metric_attributes: Optional[Dict[int, str]] = None
- self._should_prevent_trainer_and_dataloaders_deepcopy: bool = False
- # TODO: remove in 1.8
- self._running_torchscript = False
-
- self._register_sharded_tensor_state_dict_hooks_if_available()
-
- @overload
- def optimizers(self, use_pl_optimizer: Literal[True] = True) -> Union[LightningOptimizer, List[LightningOptimizer]]:
- ...
-
- @overload
- def optimizers(self, use_pl_optimizer: Literal[False]) -> Union[Optimizer, List[Optimizer]]:
- ...
-
- @overload
- def optimizers(
- self, use_pl_optimizer: bool
- ) -> Union[Optimizer, LightningOptimizer, List[Optimizer], List[LightningOptimizer]]:
- ...
-
-[docs] def optimizers(
- self, use_pl_optimizer: bool = True
- ) -> Union[Optimizer, LightningOptimizer, List[Optimizer], List[LightningOptimizer]]:
- """Returns the optimizer(s) that are being used during training. Useful for manual optimization.
-
- Args:
- use_pl_optimizer: If ``True``, will wrap the optimizer(s) in a
- :class:`~pytorch_lightning.core.optimizer.LightningOptimizer` for automatic handling of precision and
- profiling.
-
- Returns:
- A single optimizer, or a list of optimizers in case multiple ones are present.
- """
- if use_pl_optimizer:
- opts = list(self.trainer.strategy._lightning_optimizers.values())
- else:
- opts = self.trainer.optimizers
-
- # single optimizer
- if isinstance(opts, list) and len(opts) == 1 and isinstance(opts[0], (Optimizer, LightningOptimizer)):
- return opts[0]
- # multiple opts
- return opts
-
-[docs] def lr_schedulers(self) -> Optional[Union[LRSchedulerTypeUnion, List[LRSchedulerTypeUnion]]]:
- """Returns the learning rate scheduler(s) that are being used during training. Useful for manual
- optimization.
-
- Returns:
- A single scheduler, or a list of schedulers in case multiple ones are present, or ``None`` if no
- schedulers were returned in :meth:`configure_optimizers`.
- """
- if not self.trainer.lr_scheduler_configs:
- return None
-
- # ignore other keys "interval", "frequency", etc.
- lr_schedulers = [config.scheduler for config in self.trainer.lr_scheduler_configs]
-
- # single scheduler
- if len(lr_schedulers) == 1:
- return lr_schedulers[0]
-
- # multiple schedulers
- return lr_schedulers
-
- @property
- def example_input_array(self) -> Any:
- """The example input array is a specification of what the module can consume in the :meth:`forward` method.
- The return type is interpreted as follows:
-
- - Single tensor: It is assumed the model takes a single argument, i.e.,
- ``model.forward(model.example_input_array)``
- - Tuple: The input array should be interpreted as a sequence of positional arguments, i.e.,
- ``model.forward(*model.example_input_array)``
- - Dict: The input array represents named keyword arguments, i.e.,
- ``model.forward(**model.example_input_array)``
- """
- return self._example_input_array
-
- @example_input_array.setter
- def example_input_array(self, example: Any) -> None:
- self._example_input_array = example
-
- @property
- def current_epoch(self) -> int:
- """The current epoch in the ``Trainer``, or 0 if not attached."""
- return self.trainer.current_epoch if self.trainer else 0
-
- @property
- def global_step(self) -> int:
- """Total training batches seen across all epochs.
-
- If no Trainer is attached, this propery is 0.
- """
- return self.trainer.global_step if self.trainer else 0
-
- @property
- def global_rank(self) -> int:
- """The index of the current process across all nodes and devices."""
- return self.trainer.global_rank if self.trainer else 0
-
- @property
- def local_rank(self) -> int:
- """The index of the current process within a single node."""
- return self.trainer.local_rank if self.trainer else 0
-
- @property
- def on_gpu(self):
- """Returns ``True`` if this model is currently located on a GPU.
-
- Useful to set flags around the LightningModule for different CPU vs GPU behavior.
- """
- return self.device.type == "cuda"
-
- @property
- def automatic_optimization(self) -> bool:
- """If set to ``False`` you are responsible for calling ``.backward()``, ``.step()``, ``.zero_grad()``."""
- return self._automatic_optimization
-
- @automatic_optimization.setter
- def automatic_optimization(self, automatic_optimization: bool) -> None:
- self._automatic_optimization = automatic_optimization
-
- @property
- def truncated_bptt_steps(self) -> int:
- """Enables `Truncated Backpropagation Through Time` in the Trainer when set to a positive integer.
-
- It represents
- the number of times :meth:`training_step` gets called before backpropagation. If this is > 0, the
- :meth:`training_step` receives an additional argument ``hiddens`` and is expected to return a hidden state.
- """
- return self._truncated_bptt_steps
-
- @truncated_bptt_steps.setter
- def truncated_bptt_steps(self, truncated_bptt_steps: int) -> None:
- self._truncated_bptt_steps = truncated_bptt_steps
-
- @property
- def logger(self) -> Optional[LightningLoggerBase]:
- """Reference to the logger object in the Trainer."""
- return self.trainer.logger if self.trainer else None
-
- @property
- def loggers(self) -> List[LightningLoggerBase]:
- """Reference to the list of loggers in the Trainer."""
- return self.trainer.loggers if self.trainer else []
-
- def _apply_batch_transfer_handler(
- self, batch: Any, device: Optional[torch.device] = None, dataloader_idx: int = 0
- ) -> Any:
- device = device or self.device
- datahook_selector = (
- _DataHookSelector(self, None) if self.trainer is None else self.trainer._data_connector._datahook_selector
- )
-
- hook = datahook_selector.get_hook("on_before_batch_transfer")
- batch = hook(batch, dataloader_idx)
- hook = datahook_selector.get_hook("transfer_batch_to_device")
- batch = hook(batch, device, dataloader_idx)
- hook = datahook_selector.get_hook("on_after_batch_transfer")
- batch = hook(batch, dataloader_idx)
- return batch
-
-[docs] def print(self, *args, **kwargs) -> None:
- r"""
- Prints only from process 0. Use this in any distributed mode to log only once.
-
- Args:
- *args: The thing to print. The same as for Python's built-in print function.
- **kwargs: The same as for Python's built-in print function.
-
- Example::
-
- def forward(self, x):
- self.print(x, 'in forward')
-
- """
- if self.trainer.is_global_zero:
- progress_bar = self.trainer.progress_bar_callback
- if progress_bar is not None and progress_bar.is_enabled:
- progress_bar.print(*args, **kwargs)
- else:
- print(*args, **kwargs)
-
-[docs] def log(
- self,
- name: str,
- value: _METRIC_COLLECTION,
- prog_bar: bool = False,
- logger: bool = True,
- on_step: Optional[bool] = None,
- on_epoch: Optional[bool] = None,
- reduce_fx: Union[str, Callable] = "mean",
- enable_graph: bool = False,
- sync_dist: bool = False,
- sync_dist_group: Optional[Any] = None,
- add_dataloader_idx: bool = True,
- batch_size: Optional[int] = None,
- metric_attribute: Optional[str] = None,
- rank_zero_only: bool = False,
- ) -> None:
- """Log a key, value pair.
-
- Example::
-
- self.log('train_loss', loss)
-
- The default behavior per hook is documented here: :ref:`extensions/logging:Automatic Logging`.
-
- Args:
- name: key to log.
- value: value to log. Can be a ``float``, ``Tensor``, ``Metric``, or a dictionary of the former.
- prog_bar: if ``True`` logs to the progress bar.
- logger: if ``True`` logs to the logger.
- on_step: if ``True`` logs at this step. The default value is determined by the hook.
- See :ref:`extensions/logging:Automatic Logging` for details.
- on_epoch: if ``True`` logs epoch accumulated metrics. The default value is determined by the hook.
- See :ref:`extensions/logging:Automatic Logging` for details.
- reduce_fx: reduction function over step values for end of epoch. :meth:`torch.mean` by default.
- enable_graph: if ``True``, will not auto detach the graph.
- sync_dist: if ``True``, reduces the metric across devices. Use with care as this may lead to a significant
- communication overhead.
- sync_dist_group: the DDP group to sync across.
- add_dataloader_idx: if ``True``, appends the index of the current dataloader to
- the name (when using multiple dataloaders). If False, user needs to give unique names for
- each dataloader to not mix the values.
- batch_size: Current batch_size. This will be directly inferred from the loaded batch,
- but for some data structures you might need to explicitly provide it.
- metric_attribute: To restore the metric state, Lightning requires the reference of the
- :class:`torchmetrics.Metric` in your model. This is found automatically if it is a model attribute.
- rank_zero_only: Whether the value will be logged only on rank 0. This will prevent synchronization which
- would produce a deadlock as not all processes would perform this log call.
- """
- # check for invalid values
- apply_to_collection(value, dict, self.__check_not_nested, name)
- apply_to_collection(
- value, object, self.__check_allowed, name, value, wrong_dtype=(numbers.Number, Metric, Tensor, dict)
- )
-
- if self.trainer is None:
- # not an error to support testing the `*_step` methods without a `Trainer` reference
- rank_zero_warn(
- "You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet."
- " This is most likely because the model hasn't been passed to the `Trainer`"
- )
- return
- results = self.trainer._results
- if results is None:
- raise MisconfigurationException(
- "You are trying to `self.log()` but the loop's result collection is not registered"
- " yet. This is most likely because you are trying to log in a `predict` hook,"
- " but it doesn't support logging"
- )
- if self._current_fx_name is None:
- raise MisconfigurationException(
- "You are trying to `self.log()` but it is not managed by the `Trainer` control flow"
- )
-
- on_step, on_epoch = _FxValidator.check_logging_and_get_default_levels(
- self._current_fx_name, on_step=on_step, on_epoch=on_epoch
- )
-
- # make sure user doesn't introduce logic for multi-dataloaders
- if "/dataloader_idx_" in name:
- raise MisconfigurationException(
- f"You called `self.log` with the key `{name}`"
- " but it should not contain information about `dataloader_idx`"
- )
-
- value = apply_to_collection(value, numbers.Number, self.__to_tensor)
-
- if self.trainer._logger_connector.should_reset_tensors(self._current_fx_name):
- # if we started a new epoch (running its first batch) the hook name has changed
- # reset any tensors for the new hook name
- results.reset(metrics=False, fx=self._current_fx_name)
-
- if metric_attribute is None and isinstance(value, Metric):
- if self._metric_attributes is None:
- # compute once
- self._metric_attributes = {
- id(module): name for name, module in self.named_modules() if isinstance(module, Metric)
- }
- if not self._metric_attributes:
- raise MisconfigurationException(
- "Could not find the `LightningModule` attribute for the `torchmetrics.Metric` logged."
- " You can fix this by setting an attribute for the metric in your `LightningModule`."
- )
- # try to find the passed metric in the LightningModule
- metric_attribute = self._metric_attributes.get(id(value), None)
- if metric_attribute is None:
- raise MisconfigurationException(
- "Could not find the `LightningModule` attribute for the `torchmetrics.Metric` logged."
- f" You can fix this by calling `self.log({name}, ..., metric_attribute=name)` where `name` is one"
- f" of {list(self._metric_attributes.values())}"
- )
-
- if (
- self.trainer.training
- and is_param_in_hook_signature(self.training_step, "dataloader_iter", explicit=True)
- and batch_size is None
- ):
- raise MisconfigurationException(
- "With `def training_step(self, dataloader_iter)`, `self.log(..., batch_size=...)` should be provided."
- )
-
- results.log(
- self._current_fx_name,
- name,
- value,
- prog_bar=prog_bar,
- logger=logger,
- on_step=on_step,
- on_epoch=on_epoch,
- reduce_fx=reduce_fx,
- enable_graph=enable_graph,
- add_dataloader_idx=add_dataloader_idx,
- batch_size=batch_size,
- sync_dist=sync_dist and distributed_available(),
- sync_dist_fn=self.trainer.strategy.reduce or sync_ddp,
- sync_dist_group=sync_dist_group,
- metric_attribute=metric_attribute,
- rank_zero_only=rank_zero_only,
- )
-
- self.trainer._logger_connector._current_fx = self._current_fx_name
-
-[docs] def log_dict(
- self,
- dictionary: Mapping[str, _METRIC_COLLECTION],
- prog_bar: bool = False,
- logger: bool = True,
- on_step: Optional[bool] = None,
- on_epoch: Optional[bool] = None,
- reduce_fx: Union[str, Callable] = "mean",
- enable_graph: bool = False,
- sync_dist: bool = False,
- sync_dist_group: Optional[Any] = None,
- add_dataloader_idx: bool = True,
- batch_size: Optional[int] = None,
- rank_zero_only: bool = False,
- ) -> None:
- """Log a dictionary of values at once.
-
- Example::
-
- values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
- self.log_dict(values)
-
- Args:
- dictionary: key value pairs.
- The values can be a ``float``, ``Tensor``, ``Metric``, or a dictionary of the former.
- prog_bar: if ``True`` logs to the progress base.
- logger: if ``True`` logs to the logger.
- on_step: if ``True`` logs at this step.
- ``None`` auto-logs for training_step but not validation/test_step.
- The default value is determined by the hook.
- See :ref:`extensions/logging:Automatic Logging` for details.
- on_epoch: if ``True`` logs epoch accumulated metrics.
- ``None`` auto-logs for val/test step but not ``training_step``.
- The default value is determined by the hook.
- See :ref:`extensions/logging:Automatic Logging` for details.
- reduce_fx: reduction function over step values for end of epoch. :meth:`torch.mean` by default.
- enable_graph: if ``True``, will not auto-detach the graph
- sync_dist: if ``True``, reduces the metric across GPUs/TPUs. Use with care as this may lead to a significant
- communication overhead.
- sync_dist_group: the ddp group to sync across.
- add_dataloader_idx: if ``True``, appends the index of the current dataloader to
- the name (when using multiple). If ``False``, user needs to give unique names for
- each dataloader to not mix values.
- batch_size: Current batch size. This will be directly inferred from the loaded batch,
- but some data structures might need to explicitly provide it.
- rank_zero_only: Whether the value will be logged only on rank 0. This will prevent synchronization which
- would produce a deadlock as not all processes would perform this log call.
- """
- for k, v in dictionary.items():
- self.log(
- name=k,
- value=v,
- prog_bar=prog_bar,
- logger=logger,
- on_step=on_step,
- on_epoch=on_epoch,
- reduce_fx=reduce_fx,
- enable_graph=enable_graph,
- sync_dist=sync_dist,
- sync_dist_group=sync_dist_group,
- add_dataloader_idx=add_dataloader_idx,
- batch_size=batch_size,
- rank_zero_only=rank_zero_only,
- )
-
- @staticmethod
- def __check_not_nested(value: dict, name: str) -> dict:
- # self-imposed restriction. for simplicity
- if any(isinstance(v, dict) for v in value.values()):
- raise ValueError(f"`self.log({name}, {value})` was called, but nested dictionaries cannot be logged")
- return value
-
- @staticmethod
- def __check_allowed(v: Any, name: str, value: Any) -> None:
- raise ValueError(f"`self.log({name}, {value})` was called, but `{type(v).__name__}` values cannot be logged")
-
- def __to_tensor(self, value: numbers.Number) -> torch.Tensor:
- return torch.tensor(value, device=self.device)
-
- def log_grad_norm(self, grad_norm_dict: Dict[str, float]) -> None:
- """Override this method to change the default behaviour of ``log_grad_norm``.
-
- If clipping gradients, the gradients will not have been clipped yet.
-
- Args:
- grad_norm_dict: Dictionary containing current grad norm metrics
-
- Example::
-
- # DEFAULT
- def log_grad_norm(self, grad_norm_dict):
- self.log_dict(grad_norm_dict, on_step=True, on_epoch=True, prog_bar=False, logger=True)
- """
- self.log_dict(grad_norm_dict, on_step=True, on_epoch=True, prog_bar=False, logger=True)
-
-[docs] def all_gather(
- self, data: Union[torch.Tensor, Dict, List, Tuple], group: Optional[Any] = None, sync_grads: bool = False
- ):
- r"""
- Allows users to call ``self.all_gather()`` from the LightningModule, thus making the ``all_gather`` operation
- accelerator agnostic. ``all_gather`` is a function provided by accelerators to gather a tensor from several
- distributed processes.
-
- Args:
- data: int, float, tensor of shape (batch, ...), or a (possibly nested) collection thereof.
- group: the process group to gather results from. Defaults to all processes (world)
- sync_grads: flag that allows users to synchronize gradients for the all_gather operation
-
- Return:
- A tensor of shape (world_size, batch, ...), or if the input was a collection
- the output will also be a collection with tensors of this shape.
- """
- group = group if group is not None else torch.distributed.group.WORLD
- all_gather = self.trainer.strategy.all_gather
- data = convert_to_tensors(data, device=self.device)
- return apply_to_collection(data, torch.Tensor, all_gather, group=group, sync_grads=sync_grads)
-
-[docs] def forward(self, *args, **kwargs) -> Any:
- r"""
- Same as :meth:`torch.nn.Module.forward()`.
-
- Args:
- *args: Whatever you decide to pass into the forward method.
- **kwargs: Keyword arguments are also possible.
-
- Return:
- Your model's output
- """
- return super().forward(*args, **kwargs)
-
-[docs] def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
- r"""
- Here you compute and return the training loss and some additional metrics for e.g.
- the progress bar or logger.
-
- Args:
- batch (:class:`~torch.Tensor` | (:class:`~torch.Tensor`, ...) | [:class:`~torch.Tensor`, ...]):
- The output of your :class:`~torch.utils.data.DataLoader`. A tensor, tuple or list.
- batch_idx (``int``): Integer displaying index of this batch
- optimizer_idx (``int``): When using multiple optimizers, this argument will also be present.
- hiddens (``Any``): Passed in if
- :paramref:`~pytorch_lightning.core.lightning.LightningModule.truncated_bptt_steps` > 0.
-
- Return:
- Any of.
-
- - :class:`~torch.Tensor` - The loss tensor
- - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``
- - ``None`` - Training will skip to the next batch. This is only for automatic optimization.
- This is not supported for multi-GPU, TPU, IPU, or DeepSpeed.
-
- In this step you'd normally do the forward pass and calculate the loss for a batch.
- You can also do fancier things like multiple forward passes or something model specific.
-
- Example::
-
- def training_step(self, batch, batch_idx):
- x, y, z = batch
- out = self.encoder(x)
- loss = self.loss(out, x)
- return loss
-
- If you define multiple optimizers, this step will be called with an additional
- ``optimizer_idx`` parameter.
-
- .. code-block:: python
-
- # Multiple optimizers (e.g.: GANs)
- def training_step(self, batch, batch_idx, optimizer_idx):
- if optimizer_idx == 0:
- # do training_step with encoder
- ...
- if optimizer_idx == 1:
- # do training_step with decoder
- ...
-
-
- If you add truncated back propagation through time you will also get an additional
- argument with the hidden states of the previous step.
-
- .. code-block:: python
-
- # Truncated back-propagation through time
- def training_step(self, batch, batch_idx, hiddens):
- # hiddens are the hidden states from the previous truncated backprop step
- out, hiddens = self.lstm(data, hiddens)
- loss = ...
- return {"loss": loss, "hiddens": hiddens}
-
- Note:
- The loss value shown in the progress bar is smoothed (averaged) over the last values,
- so it differs from the actual loss returned in train/validation step.
- """
- rank_zero_warn("`training_step` must be implemented to be used with the Lightning Trainer")
-
-[docs] def training_step_end(self, step_output: STEP_OUTPUT) -> STEP_OUTPUT:
- """Use this when training with dp or ddp2 because :meth:`training_step` will operate on only part of the
- batch. However, this is still optional and only needed for things like softmax or NCE loss.
-
- Note:
- If you later switch to ddp or some other mode, this will still be called
- so that you don't have to change your code
-
- .. code-block:: python
-
- # pseudocode
- sub_batches = split_batches_for_dp(batch)
- step_output = [training_step(sub_batch) for sub_batch in sub_batches]
- training_step_end(step_output)
-
- Args:
- step_output: What you return in `training_step` for each batch part.
-
- Return:
- Anything
-
- When using dp/ddp2 distributed backends, only a portion of the batch is inside the training_step:
-
- .. code-block:: python
-
- def training_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self(x)
-
- # softmax uses only a portion of the batch in the denominator
- loss = self.softmax(out)
- loss = nce_loss(loss)
- return loss
-
- If you wish to do something with all the parts of the batch, then use this method to do it:
-
- .. code-block:: python
-
- def training_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self.encoder(x)
- return {"pred": out}
-
-
- def training_step_end(self, training_step_outputs):
- gpu_0_pred = training_step_outputs[0]["pred"]
- gpu_1_pred = training_step_outputs[1]["pred"]
- gpu_n_pred = training_step_outputs[n]["pred"]
-
- # this softmax now uses the full batch
- loss = nce_loss([gpu_0_pred, gpu_1_pred, gpu_n_pred])
- return loss
-
- See Also:
- See the :ref:`accelerators/gpu:Multi GPU Training` guide for more details.
- """
-
-[docs] def training_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
- """Called at the end of the training epoch with the outputs of all training steps. Use this in case you
- need to do something with all the outputs returned by :meth:`training_step`.
-
- .. code-block:: python
-
- # the pseudocode for these calls
- train_outs = []
- for train_batch in train_data:
- out = training_step(train_batch)
- train_outs.append(out)
- training_epoch_end(train_outs)
-
- Args:
- outputs: List of outputs you defined in :meth:`training_step`. If there are multiple optimizers or when
- using ``truncated_bptt_steps > 0``, the lists have the dimensions
- (n_batches, tbptt_steps, n_optimizers). Dimensions of length 1 are squeezed.
-
- Return:
- None
-
- Note:
- If this method is not overridden, this won't be called.
-
- .. code-block:: python
-
- def training_epoch_end(self, training_step_outputs):
- # do something with all training_step outputs
- for out in training_step_outputs:
- ...
- """
-
-[docs] def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
- r"""
- Operates on a single batch of data from the validation set.
- In this step you'd might generate examples or calculate anything of interest like accuracy.
-
- .. code-block:: python
-
- # the pseudocode for these calls
- val_outs = []
- for val_batch in val_data:
- out = validation_step(val_batch)
- val_outs.append(out)
- validation_epoch_end(val_outs)
-
- Args:
- batch: The output of your :class:`~torch.utils.data.DataLoader`.
- batch_idx: The index of this batch.
- dataloader_idx: The index of the dataloader that produced this batch.
- (only if multiple val dataloaders used)
-
- Return:
- - Any object or value
- - ``None`` - Validation will skip to the next batch
-
- .. code-block:: python
-
- # pseudocode of order
- val_outs = []
- for val_batch in val_data:
- out = validation_step(val_batch)
- if defined("validation_step_end"):
- out = validation_step_end(out)
- val_outs.append(out)
- val_outs = validation_epoch_end(val_outs)
-
-
- .. code-block:: python
-
- # if you have one val dataloader:
- def validation_step(self, batch, batch_idx):
- ...
-
-
- # if you have multiple val dataloaders:
- def validation_step(self, batch, batch_idx, dataloader_idx=0):
- ...
-
- Examples::
-
- # CASE 1: A single validation dataset
- def validation_step(self, batch, batch_idx):
- x, y = batch
-
- # implement your own
- out = self(x)
- loss = self.loss(out, y)
-
- # log 6 example images
- # or generated text... or whatever
- sample_imgs = x[:6]
- grid = torchvision.utils.make_grid(sample_imgs)
- self.logger.experiment.add_image('example_images', grid, 0)
-
- # calculate acc
- labels_hat = torch.argmax(out, dim=1)
- val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
- # log the outputs!
- self.log_dict({'val_loss': loss, 'val_acc': val_acc})
-
- If you pass in multiple val dataloaders, :meth:`validation_step` will have an additional argument. We recommend
- setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.
-
- .. code-block:: python
-
- # CASE 2: multiple validation dataloaders
- def validation_step(self, batch, batch_idx, dataloader_idx=0):
- # dataloader_idx tells you which dataset this is.
- ...
-
- Note:
- If you don't need to validate you don't need to implement this method.
-
- Note:
- When the :meth:`validation_step` is called, the model has been put in eval mode
- and PyTorch gradients have been disabled. At the end of validation,
- the model goes back to training mode and gradients are enabled.
- """
-
-[docs] def validation_step_end(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
- """Use this when validating with dp or ddp2 because :meth:`validation_step` will operate on only part of
- the batch. However, this is still optional and only needed for things like softmax or NCE loss.
-
- Note:
- If you later switch to ddp or some other mode, this will still be called
- so that you don't have to change your code.
-
- .. code-block:: python
-
- # pseudocode
- sub_batches = split_batches_for_dp(batch)
- step_output = [validation_step(sub_batch) for sub_batch in sub_batches]
- validation_step_end(step_output)
-
- Args:
- step_output: What you return in :meth:`validation_step` for each batch part.
-
- Return:
- None or anything
-
- .. code-block:: python
-
- # WITHOUT validation_step_end
- # if used in DP or DDP2, this batch is 1/num_gpus large
- def validation_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self.encoder(x)
- loss = self.softmax(out)
- loss = nce_loss(loss)
- self.log("val_loss", loss)
-
-
- # --------------
- # with validation_step_end to do softmax over the full batch
- def validation_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self(x)
- return out
-
-
- def validation_step_end(self, val_step_outputs):
- for out in val_step_outputs:
- ...
-
- See Also:
- See the :ref:`accelerators/gpu:Multi GPU Training` guide for more details.
- """
-
-[docs] def validation_epoch_end(self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]) -> None:
- """Called at the end of the validation epoch with the outputs of all validation steps.
-
- .. code-block:: python
-
- # the pseudocode for these calls
- val_outs = []
- for val_batch in val_data:
- out = validation_step(val_batch)
- val_outs.append(out)
- validation_epoch_end(val_outs)
-
- Args:
- outputs: List of outputs you defined in :meth:`validation_step`, or if there
- are multiple dataloaders, a list containing a list of outputs for each dataloader.
-
- Return:
- None
-
- Note:
- If you didn't define a :meth:`validation_step`, this won't be called.
-
- Examples:
- With a single dataloader:
-
- .. code-block:: python
-
- def validation_epoch_end(self, val_step_outputs):
- for out in val_step_outputs:
- ...
-
- With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
- one entry per dataloader, while the inner list contains the individual outputs of
- each validation step for that dataloader.
-
- .. code-block:: python
-
- def validation_epoch_end(self, outputs):
- for dataloader_output_result in outputs:
- dataloader_outs = dataloader_output_result.dataloader_i_outputs
-
- self.log("final_metric", final_value)
- """
-
-[docs] def test_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
- r"""
- Operates on a single batch of data from the test set.
- In this step you'd normally generate examples or calculate anything of interest
- such as accuracy.
-
- .. code-block:: python
-
- # the pseudocode for these calls
- test_outs = []
- for test_batch in test_data:
- out = test_step(test_batch)
- test_outs.append(out)
- test_epoch_end(test_outs)
-
- Args:
- batch: The output of your :class:`~torch.utils.data.DataLoader`.
- batch_idx: The index of this batch.
- dataloader_id: The index of the dataloader that produced this batch.
- (only if multiple test dataloaders used).
-
- Return:
- Any of.
-
- - Any object or value
- - ``None`` - Testing will skip to the next batch
-
- .. code-block:: python
-
- # if you have one test dataloader:
- def test_step(self, batch, batch_idx):
- ...
-
-
- # if you have multiple test dataloaders:
- def test_step(self, batch, batch_idx, dataloader_idx=0):
- ...
-
- Examples::
-
- # CASE 1: A single test dataset
- def test_step(self, batch, batch_idx):
- x, y = batch
-
- # implement your own
- out = self(x)
- loss = self.loss(out, y)
-
- # log 6 example images
- # or generated text... or whatever
- sample_imgs = x[:6]
- grid = torchvision.utils.make_grid(sample_imgs)
- self.logger.experiment.add_image('example_images', grid, 0)
-
- # calculate acc
- labels_hat = torch.argmax(out, dim=1)
- test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
- # log the outputs!
- self.log_dict({'test_loss': loss, 'test_acc': test_acc})
-
- If you pass in multiple test dataloaders, :meth:`test_step` will have an additional argument. We recommend
- setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.
-
- .. code-block:: python
-
- # CASE 2: multiple test dataloaders
- def test_step(self, batch, batch_idx, dataloader_idx=0):
- # dataloader_idx tells you which dataset this is.
- ...
-
- Note:
- If you don't need to test you don't need to implement this method.
-
- Note:
- When the :meth:`test_step` is called, the model has been put in eval mode and
- PyTorch gradients have been disabled. At the end of the test epoch, the model goes back
- to training mode and gradients are enabled.
- """
-
-[docs] def test_step_end(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
- """Use this when testing with dp or ddp2 because :meth:`test_step` will operate on only part of the batch.
- However, this is still optional and only needed for things like softmax or NCE loss.
-
- Note:
- If you later switch to ddp or some other mode, this will still be called
- so that you don't have to change your code.
-
- .. code-block:: python
-
- # pseudocode
- sub_batches = split_batches_for_dp(batch)
- step_output = [test_step(sub_batch) for sub_batch in sub_batches]
- test_step_end(step_output)
-
- Args:
- step_output: What you return in :meth:`test_step` for each batch part.
-
- Return:
- None or anything
-
- .. code-block:: python
-
- # WITHOUT test_step_end
- # if used in DP or DDP2, this batch is 1/num_gpus large
- def test_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self(x)
- loss = self.softmax(out)
- self.log("test_loss", loss)
-
-
- # --------------
- # with test_step_end to do softmax over the full batch
- def test_step(self, batch, batch_idx):
- # batch is 1/num_gpus big
- x, y = batch
-
- out = self.encoder(x)
- return out
-
-
- def test_step_end(self, output_results):
- # this out is now the full size of the batch
- all_test_step_outs = output_results.out
- loss = nce_loss(all_test_step_outs)
- self.log("test_loss", loss)
-
- See Also:
- See the :ref:`accelerators/gpu:Multi GPU Training` guide for more details.
- """
-
-[docs] def test_epoch_end(self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]) -> None:
- """Called at the end of a test epoch with the output of all test steps.
-
- .. code-block:: python
-
- # the pseudocode for these calls
- test_outs = []
- for test_batch in test_data:
- out = test_step(test_batch)
- test_outs.append(out)
- test_epoch_end(test_outs)
-
- Args:
- outputs: List of outputs you defined in :meth:`test_step_end`, or if there
- are multiple dataloaders, a list containing a list of outputs for each dataloader
-
- Return:
- None
-
- Note:
- If you didn't define a :meth:`test_step`, this won't be called.
-
- Examples:
- With a single dataloader:
-
- .. code-block:: python
-
- def test_epoch_end(self, outputs):
- # do something with the outputs of all test batches
- all_test_preds = test_step_outputs.predictions
-
- some_result = calc_all_results(all_test_preds)
- self.log(some_result)
-
- With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
- one entry per dataloader, while the inner list contains the individual outputs of
- each test step for that dataloader.
-
- .. code-block:: python
-
- def test_epoch_end(self, outputs):
- final_value = 0
- for dataloader_outputs in outputs:
- for test_step_out in dataloader_outputs:
- # do something
- final_value += test_step_out
-
- self.log("final_metric", final_value)
- """
-
-[docs] def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
- """Step function called during :meth:`~pytorch_lightning.trainer.trainer.Trainer.predict`. By default, it
- calls :meth:`~pytorch_lightning.core.lightning.LightningModule.forward`. Override to add any processing
- logic.
-
- The :meth:`~pytorch_lightning.core.lightning.LightningModule.predict_step` is used
- to scale inference on multi-devices.
-
- To prevent an OOM error, it is possible to use :class:`~pytorch_lightning.callbacks.BasePredictionWriter`
- callback to write the predictions to disk or database after each batch or on epoch end.
-
- The :class:`~pytorch_lightning.callbacks.BasePredictionWriter` should be used while using a spawn
- based accelerator. This happens for ``Trainer(strategy="ddp_spawn")``
- or training on 8 TPU cores with ``Trainer(accelerator="tpu", devices=8)`` as predictions won't be returned.
-
- Example ::
-
- class MyModel(LightningModule):
-
- def predicts_step(self, batch, batch_idx, dataloader_idx=0):
- return self(batch)
-
- dm = ...
- model = MyModel()
- trainer = Trainer(accelerator="gpu", devices=2)
- predictions = trainer.predict(model, dm)
-
-
- Args:
- batch: Current batch.
- batch_idx: Index of current batch.
- dataloader_idx: Index of the current dataloader.
-
- Return:
- Predicted output
- """
- return self(batch)
-
-[docs] def configure_callbacks(self) -> Union[Sequence[Callback], Callback]:
- """Configure model-specific callbacks. When the model gets attached, e.g., when ``.fit()`` or ``.test()``
- gets called, the list or a callback returned here will be merged with the list of callbacks passed to the
- Trainer's ``callbacks`` argument. If a callback returned here has the same type as one or several callbacks
- already present in the Trainer's callbacks list, it will take priority and replace them. In addition,
- Lightning will make sure :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint` callbacks
- run last.
-
- Return:
- A callback or a list of callbacks which will extend the list of callbacks in the Trainer.
-
- Example::
-
- def configure_callbacks(self):
- early_stop = EarlyStopping(monitor="val_acc", mode="max")
- checkpoint = ModelCheckpoint(monitor="val_loss")
- return [early_stop, checkpoint]
-
- Note:
- Certain callback methods like :meth:`~pytorch_lightning.callbacks.base.Callback.on_init_start`
- will never be invoked on the new callbacks returned here.
- """
- return []
-
-[docs] def configure_optimizers(self):
- r"""
- Choose what optimizers and learning-rate schedulers to use in your optimization.
- Normally you'd need one. But in the case of GANs or similar you might have multiple.
-
- Return:
- Any of these 6 options.
-
- - **Single optimizer**.
- - **List or Tuple** of optimizers.
- - **Two lists** - The first list has multiple optimizers, and the second has multiple LR schedulers
- (or multiple ``lr_scheduler_config``).
- - **Dictionary**, with an ``"optimizer"`` key, and (optionally) a ``"lr_scheduler"``
- key whose value is a single LR scheduler or ``lr_scheduler_config``.
- - **Tuple of dictionaries** as described above, with an optional ``"frequency"`` key.
- - **None** - Fit will run without any optimizer.
-
- The ``lr_scheduler_config`` is a dictionary which contains the scheduler and its associated configuration.
- The default configuration is shown below.
-
- .. code-block:: python
-
- lr_scheduler_config = {
- # REQUIRED: The scheduler instance
- "scheduler": lr_scheduler,
- # The unit of the scheduler's step size, could also be 'step'.
- # 'epoch' updates the scheduler on epoch end whereas 'step'
- # updates it after a optimizer update.
- "interval": "epoch",
- # How many epochs/steps should pass between calls to
- # `scheduler.step()`. 1 corresponds to updating the learning
- # rate after every epoch/step.
- "frequency": 1,
- # Metric to to monitor for schedulers like `ReduceLROnPlateau`
- "monitor": "val_loss",
- # If set to `True`, will enforce that the value specified 'monitor'
- # is available when the scheduler is updated, thus stopping
- # training if not found. If set to `False`, it will only produce a warning
- "strict": True,
- # If using the `LearningRateMonitor` callback to monitor the
- # learning rate progress, this keyword can be used to specify
- # a custom logged name
- "name": None,
- }
-
- When there are schedulers in which the ``.step()`` method is conditioned on a value, such as the
- :class:`torch.optim.lr_scheduler.ReduceLROnPlateau` scheduler, Lightning requires that the
- ``lr_scheduler_config`` contains the keyword ``"monitor"`` set to the metric name that the scheduler
- should be conditioned on.
-
- .. testcode::
-
- # The ReduceLROnPlateau scheduler requires a monitor
- def configure_optimizers(self):
- optimizer = Adam(...)
- return {
- "optimizer": optimizer,
- "lr_scheduler": {
- "scheduler": ReduceLROnPlateau(optimizer, ...),
- "monitor": "metric_to_track",
- "frequency": "indicates how often the metric is updated"
- # If "monitor" references validation metrics, then "frequency" should be set to a
- # multiple of "trainer.check_val_every_n_epoch".
- },
- }
-
-
- # In the case of two optimizers, only one using the ReduceLROnPlateau scheduler
- def configure_optimizers(self):
- optimizer1 = Adam(...)
- optimizer2 = SGD(...)
- scheduler1 = ReduceLROnPlateau(optimizer1, ...)
- scheduler2 = LambdaLR(optimizer2, ...)
- return (
- {
- "optimizer": optimizer1,
- "lr_scheduler": {
- "scheduler": scheduler1,
- "monitor": "metric_to_track",
- },
- },
- {"optimizer": optimizer2, "lr_scheduler": scheduler2},
- )
-
- Metrics can be made available to monitor by simply logging it using
- ``self.log('metric_to_track', metric_val)`` in your :class:`~pytorch_lightning.core.lightning.LightningModule`.
-
- Note:
- The ``frequency`` value specified in a dict along with the ``optimizer`` key is an int corresponding
- to the number of sequential batches optimized with the specific optimizer.
- It should be given to none or to all of the optimizers.
- There is a difference between passing multiple optimizers in a list,
- and passing multiple optimizers in dictionaries with a frequency of 1:
-
- - In the former case, all optimizers will operate on the given batch in each optimization step.
- - In the latter, only one optimizer will operate on the given batch at every step.
-
- This is different from the ``frequency`` value specified in the ``lr_scheduler_config`` mentioned above.
-
- .. code-block:: python
-
- def configure_optimizers(self):
- optimizer_one = torch.optim.SGD(self.model.parameters(), lr=0.01)
- optimizer_two = torch.optim.SGD(self.model.parameters(), lr=0.01)
- return [
- {"optimizer": optimizer_one, "frequency": 5},
- {"optimizer": optimizer_two, "frequency": 10},
- ]
-
- In this example, the first optimizer will be used for the first 5 steps,
- the second optimizer for the next 10 steps and that cycle will continue.
- If an LR scheduler is specified for an optimizer using the ``lr_scheduler`` key in the above dict,
- the scheduler will only be updated when its optimizer is being used.
-
- Examples::
-
- # most cases. no learning rate scheduler
- def configure_optimizers(self):
- return Adam(self.parameters(), lr=1e-3)
-
- # multiple optimizer case (e.g.: GAN)
- def configure_optimizers(self):
- gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
- dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
- return gen_opt, dis_opt
-
- # example with learning rate schedulers
- def configure_optimizers(self):
- gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
- dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
- dis_sch = CosineAnnealing(dis_opt, T_max=10)
- return [gen_opt, dis_opt], [dis_sch]
-
- # example with step-based learning rate schedulers
- # each optimizer has its own scheduler
- def configure_optimizers(self):
- gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
- dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
- gen_sch = {
- 'scheduler': ExponentialLR(gen_opt, 0.99),
- 'interval': 'step' # called after each training step
- }
- dis_sch = CosineAnnealing(dis_opt, T_max=10) # called every epoch
- return [gen_opt, dis_opt], [gen_sch, dis_sch]
-
- # example with optimizer frequencies
- # see training procedure in `Improved Training of Wasserstein GANs`, Algorithm 1
- # https://arxiv.org/abs/1704.00028
- def configure_optimizers(self):
- gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
- dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
- n_critic = 5
- return (
- {'optimizer': dis_opt, 'frequency': n_critic},
- {'optimizer': gen_opt, 'frequency': 1}
- )
-
- Note:
- Some things to know:
-
- - Lightning calls ``.backward()`` and ``.step()`` on each optimizer and learning rate scheduler as needed.
- - If you use 16-bit precision (``precision=16``), Lightning will automatically handle the optimizers.
- - If you use multiple optimizers, :meth:`training_step` will have an additional ``optimizer_idx`` parameter.
- - If you use :class:`torch.optim.LBFGS`, Lightning handles the closure function automatically for you.
- - If you use multiple optimizers, gradients will be calculated only for the parameters of current optimizer
- at each training step.
- - If you need to control how often those optimizers step or override the default ``.step()`` schedule,
- override the :meth:`optimizer_step` hook.
- """
- rank_zero_warn("`configure_optimizers` must be implemented to be used with the Lightning Trainer")
-
-[docs] def manual_backward(self, loss: Tensor, *args, **kwargs) -> None:
- """Call this directly from your :meth:`training_step` when doing optimizations manually. By using this,
- Lightning can ensure that all the proper scaling gets applied when using mixed precision.
-
- See :ref:`manual optimization<common/optimization:Manual optimization>` for more examples.
-
- Example::
-
- def training_step(...):
- opt = self.optimizers()
- loss = ...
- opt.zero_grad()
- # automatically applies scaling, etc...
- self.manual_backward(loss)
- opt.step()
-
- Args:
- loss: The tensor on which to compute gradients. Must have a graph attached.
- *args: Additional positional arguments to be forwarded to :meth:`~torch.Tensor.backward`
- **kwargs: Additional keyword arguments to be forwarded to :meth:`~torch.Tensor.backward`
- """
- self._verify_is_manual_optimization("manual_backward")
- self.trainer.strategy.backward(loss, None, None, *args, **kwargs)
-
-[docs] def backward(
- self, loss: Tensor, optimizer: Optional[Optimizer], optimizer_idx: Optional[int], *args, **kwargs
- ) -> None:
- """Called to perform backward on the loss returned in :meth:`training_step`. Override this hook with your
- own implementation if you need to.
-
- Args:
- loss: The loss tensor returned by :meth:`training_step`. If gradient accumulation is used, the loss here
- holds the normalized value (scaled by 1 / accumulation steps).
- optimizer: Current optimizer being used. ``None`` if using manual optimization.
- optimizer_idx: Index of the current optimizer being used. ``None`` if using manual optimization.
-
- Example::
-
- def backward(self, loss, optimizer, optimizer_idx):
- loss.backward()
- """
- loss.backward(*args, **kwargs)
-
-[docs] def toggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer], optimizer_idx: int) -> None:
- """Makes sure only the gradients of the current optimizer's parameters are calculated in the training step
- to prevent dangling gradients in multiple-optimizer setup.
-
- This is only called automatically when automatic optimization is enabled and multiple optimizers are used.
- It works with :meth:`untoggle_optimizer` to make sure ``param_requires_grad_state`` is properly reset.
-
- Args:
- optimizer: The optimizer to toggle.
- optimizer_idx: The index of the optimizer to toggle.
- """
- # Iterate over all optimizer parameters to preserve their `requires_grad` information
- # in case these are pre-defined during `configure_optimizers`
- param_requires_grad_state = {}
- for opt in self.trainer.optimizers:
- for group in opt.param_groups:
- for param in group["params"]:
- # If a param already appear in param_requires_grad_state, continue
- if param in param_requires_grad_state:
- continue
- param_requires_grad_state[param] = param.requires_grad
- param.requires_grad = False
-
- # Then iterate over the current optimizer's parameters and set its `requires_grad`
- # properties accordingly
- for group in optimizer.param_groups:
- for param in group["params"]:
- param.requires_grad = param_requires_grad_state[param]
- self._param_requires_grad_state = param_requires_grad_state
-
-[docs] def untoggle_optimizer(self, optimizer_idx: int) -> None:
- """Resets the state of required gradients that were toggled with :meth:`toggle_optimizer`.
-
- This is only called automatically when automatic optimization is enabled and multiple optimizers are used.
-
- Args:
- optimizer_idx: The index of the optimizer to untoggle.
- """
- for opt_idx, opt in enumerate(self.trainer.optimizers):
- if optimizer_idx != opt_idx:
- for group in opt.param_groups:
- for param in group["params"]:
- if param in self._param_requires_grad_state:
- param.requires_grad = self._param_requires_grad_state[param]
- # save memory
- self._param_requires_grad_state = {}
-
- def clip_gradients(
- self,
- optimizer: Optimizer,
- gradient_clip_val: Optional[Union[int, float]] = None,
- gradient_clip_algorithm: Optional[str] = None,
- ):
- """Handles gradient clipping internally.
-
- Note:
- Do not override this method. If you want to customize gradient clipping, consider
- using :meth:`configure_gradient_clipping` method.
-
- Args:
- optimizer: Current optimizer being used.
- gradient_clip_val: The value at which to clip gradients.
- gradient_clip_algorithm: The gradient clipping algorithm to use. Pass ``gradient_clip_algorithm="value"``
- to clip by value, and ``gradient_clip_algorithm="norm"`` to clip by norm.
- """
- if gradient_clip_val is None:
- gradient_clip_val = self.trainer.gradient_clip_val or 0.0
- elif self.trainer.gradient_clip_val is not None and self.trainer.gradient_clip_val != gradient_clip_val:
- raise MisconfigurationException(
- f"You have set `Trainer(gradient_clip_val={self.trainer.gradient_clip_val!r})`"
- f" and have passed `clip_gradients(gradient_clip_val={gradient_clip_val!r})`."
- " Please use only one of them."
- )
-
- if gradient_clip_algorithm is None:
- gradient_clip_algorithm = self.trainer.gradient_clip_algorithm or "norm"
- else:
- gradient_clip_algorithm = gradient_clip_algorithm.lower()
- if (
- self.trainer.gradient_clip_algorithm is not None
- and self.trainer.gradient_clip_algorithm != gradient_clip_algorithm
- ):
- raise MisconfigurationException(
- f"You have set `Trainer(gradient_clip_algorithm={self.trainer.gradient_clip_algorithm.value!r})`"
- f" and have passed `clip_gradients(gradient_clip_algorithm={gradient_clip_algorithm!r})"
- " Please use only one of them."
- )
-
- if not isinstance(gradient_clip_val, (int, float)):
- raise TypeError(f"`gradient_clip_val` should be an int or a float. Got {gradient_clip_val}.")
-
- if not GradClipAlgorithmType.supported_type(gradient_clip_algorithm.lower()):
- raise MisconfigurationException(
- f"`gradient_clip_algorithm` {gradient_clip_algorithm} is invalid."
- f" Allowed algorithms: {GradClipAlgorithmType.supported_types()}."
- )
-
- gradient_clip_algorithm = GradClipAlgorithmType(gradient_clip_algorithm)
- self.trainer.precision_plugin.clip_gradients(optimizer, gradient_clip_val, gradient_clip_algorithm)
-
-[docs] def configure_gradient_clipping(
- self,
- optimizer: Optimizer,
- optimizer_idx: int,
- gradient_clip_val: Optional[Union[int, float]] = None,
- gradient_clip_algorithm: Optional[str] = None,
- ):
- """Perform gradient clipping for the optimizer parameters. Called before :meth:`optimizer_step`.
-
- Args:
- optimizer: Current optimizer being used.
- optimizer_idx: Index of the current optimizer being used.
- gradient_clip_val: The value at which to clip gradients. By default value passed in Trainer
- will be available here.
- gradient_clip_algorithm: The gradient clipping algorithm to use. By default value
- passed in Trainer will be available here.
-
- Example::
-
- # Perform gradient clipping on gradients associated with discriminator (optimizer_idx=1) in GAN
- def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_val, gradient_clip_algorithm):
- if optimizer_idx == 1:
- # Lightning will handle the gradient clipping
- self.clip_gradients(
- optimizer,
- gradient_clip_val=gradient_clip_val,
- gradient_clip_algorithm=gradient_clip_algorithm
- )
- else:
- # implement your own custom logic to clip gradients for generator (optimizer_idx=0)
- """
- self.clip_gradients(
- optimizer, gradient_clip_val=gradient_clip_val, gradient_clip_algorithm=gradient_clip_algorithm
- )
-
- def lr_scheduler_step(
- self,
- scheduler: LRSchedulerTypeUnion,
- optimizer_idx: int,
- metric: Optional[Any],
- ) -> None:
- r"""
- Override this method to adjust the default way the
- :class:`~pytorch_lightning.trainer.trainer.Trainer` calls each scheduler.
- By default, Lightning calls ``step()`` and as shown in the example
- for each scheduler based on its ``interval``.
-
- Args:
- scheduler: Learning rate scheduler.
- optimizer_idx: Index of the optimizer associated with this scheduler.
- metric: Value of the monitor used for schedulers like ``ReduceLROnPlateau``.
-
- Examples::
-
- # DEFAULT
- def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
- if metric is None:
- scheduler.step()
- else:
- scheduler.step(metric)
-
- # Alternative way to update schedulers if it requires an epoch value
- def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
- scheduler.step(epoch=self.current_epoch)
-
- """
- if metric is None:
- scheduler.step()
- else:
- scheduler.step(metric)
-
-[docs] def optimizer_step(
- self,
- epoch: int,
- batch_idx: int,
- optimizer: Union[Optimizer, LightningOptimizer],
- optimizer_idx: int = 0,
- optimizer_closure: Optional[Callable[[], Any]] = None,
- on_tpu: bool = False,
- using_native_amp: bool = False,
- using_lbfgs: bool = False,
- ) -> None:
- r"""
- Override this method to adjust the default way the :class:`~pytorch_lightning.trainer.trainer.Trainer` calls
- each optimizer.
-
- By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example once per optimizer.
- This method (and ``zero_grad()``) won't be called during the accumulation phase when
- ``Trainer(accumulate_grad_batches != 1)``. Overriding this hook has no benefit with manual optimization.
-
- Args:
- epoch: Current epoch
- batch_idx: Index of current batch
- optimizer: A PyTorch optimizer
- optimizer_idx: If you used multiple optimizers, this indexes into that list.
- optimizer_closure: The optimizer closure. This closure must be executed as it includes the
- calls to ``training_step()``, ``optimizer.zero_grad()``, and ``backward()``.
- on_tpu: ``True`` if TPU backward is required
- using_native_amp: ``True`` if using native amp
- using_lbfgs: True if the matching optimizer is :class:`torch.optim.LBFGS`
-
- Examples::
-
- # DEFAULT
- def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx,
- optimizer_closure, on_tpu, using_native_amp, using_lbfgs):
- optimizer.step(closure=optimizer_closure)
-
- # Alternating schedule for optimizer steps (i.e.: GANs)
- def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx,
- optimizer_closure, on_tpu, using_native_amp, using_lbfgs):
- # update generator opt every step
- if optimizer_idx == 0:
- optimizer.step(closure=optimizer_closure)
-
- # update discriminator opt every 2 steps
- if optimizer_idx == 1:
- if (batch_idx + 1) % 2 == 0 :
- optimizer.step(closure=optimizer_closure)
- else:
- # call the closure by itself to run `training_step` + `backward` without an optimizer step
- optimizer_closure()
-
- # ...
- # add as many optimizers as you want
-
- Here's another example showing how to use this for more advanced things such as
- learning rate warm-up:
-
- .. code-block:: python
-
- # learning rate warm-up
- def optimizer_step(
- self,
- epoch,
- batch_idx,
- optimizer,
- optimizer_idx,
- optimizer_closure,
- on_tpu,
- using_native_amp,
- using_lbfgs,
- ):
- # update params
- optimizer.step(closure=optimizer_closure)
-
- # manually warm up lr without a scheduler
- if self.trainer.global_step < 500:
- lr_scale = min(1.0, float(self.trainer.global_step + 1) / 500.0)
- for pg in optimizer.param_groups:
- pg["lr"] = lr_scale * self.learning_rate
-
- """
- optimizer.step(closure=optimizer_closure)
-
-[docs] def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int):
- """Override this method to change the default behaviour of ``optimizer.zero_grad()``.
-
- Args:
- epoch: Current epoch
- batch_idx: Index of current batch
- optimizer: A PyTorch optimizer
- optimizer_idx: If you used multiple optimizers this indexes into that list.
-
- Examples::
-
- # DEFAULT
- def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
- optimizer.zero_grad()
-
- # Set gradients to `None` instead of zero to improve performance.
- def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
- optimizer.zero_grad(set_to_none=True)
-
- See :meth:`torch.optim.Optimizer.zero_grad` for the explanation of the above example.
- """
- optimizer.zero_grad()
-
-[docs] def tbptt_split_batch(self, batch: Any, split_size: int) -> List[Any]:
- r"""
- When using truncated backpropagation through time, each batch must be split along the
- time dimension. Lightning handles this by default, but for custom behavior override
- this function.
-
- Args:
- batch: Current batch
- split_size: The size of the split
-
- Return:
- List of batch splits. Each split will be passed to :meth:`training_step` to enable truncated
- back propagation through time. The default implementation splits root level Tensors and
- Sequences at dim=1 (i.e. time dim). It assumes that each time dim is the same length.
-
- Examples::
-
- def tbptt_split_batch(self, batch, split_size):
- splits = []
- for t in range(0, time_dims[0], split_size):
- batch_split = []
- for i, x in enumerate(batch):
- if isinstance(x, torch.Tensor):
- split_x = x[:, t:t + split_size]
- elif isinstance(x, collections.Sequence):
- split_x = [None] * len(x)
- for batch_idx in range(len(x)):
- split_x[batch_idx] = x[batch_idx][t:t + split_size]
- batch_split.append(split_x)
- splits.append(batch_split)
- return splits
-
- Note:
- Called in the training loop after
- :meth:`~pytorch_lightning.callbacks.base.Callback.on_train_batch_start`
- if :paramref:`~pytorch_lightning.core.lightning.LightningModule.truncated_bptt_steps` > 0.
- Each returned batch split is passed separately to :meth:`training_step`.
- """
- time_dims = [len(x[0]) for x in batch if isinstance(x, (torch.Tensor, collections.Sequence))]
- assert len(time_dims) >= 1, "Unable to determine batch time dimension"
- assert all(x == time_dims[0] for x in time_dims), "Batch time dimension length is ambiguous"
-
- splits = []
- for t in range(0, time_dims[0], split_size):
- batch_split = []
- for i, x in enumerate(batch):
- if isinstance(x, torch.Tensor):
- split_x = x[:, t : t + split_size]
- elif isinstance(x, collections.Sequence):
- split_x = [None] * len(x)
- for batch_idx in range(len(x)):
- split_x[batch_idx] = x[batch_idx][t : t + split_size]
-
- batch_split.append(split_x)
-
- splits.append(batch_split)
-
- return splits
-
- def summarize(self, max_depth: int = 1) -> ModelSummary:
- """Summarize this LightningModule.
-
- .. deprecated:: v1.5
- This method was deprecated in v1.5 in favor of `pytorch_lightning.utilities.model_summary.summarize`
- and will be removed in v1.7.
-
- Args:
- max_depth: The maximum depth of layer nesting that the summary will include. A value of 0 turns the
- layer summary off. Default: 1.
-
- Return:
- The model summary object
- """
- rank_zero_deprecation(
- "The `LightningModule.summarize` method is deprecated in v1.5 and will be removed in v1.7. "
- "Use `pytorch_lightning.utilities.model_summary.summarize` instead.",
- stacklevel=6,
- )
-
- return summarize(self, max_depth)
-
-[docs] def freeze(self) -> None:
- r"""
- Freeze all params for inference.
-
- Example::
-
- model = MyLightningModule(...)
- model.freeze()
-
- """
- for param in self.parameters():
- param.requires_grad = False
-
- self.eval()
-
-[docs] def unfreeze(self) -> None:
- """Unfreeze all parameters for training.
-
- .. code-block:: python
-
- model = MyLightningModule(...)
- model.unfreeze()
- """
- for param in self.parameters():
- param.requires_grad = True
-
- self.train()
-
- def get_progress_bar_dict(self) -> Dict[str, Union[int, str]]:
- r"""
- .. deprecated:: v1.5
- This method was deprecated in v1.5 in favor of
- `pytorch_lightning.callbacks.progress.base.get_metrics` and will be removed in v1.7.
-
- Implement this to override the default items displayed in the progress bar.
- By default it includes the average loss value, split index of BPTT (if used)
- and the version of the experiment when using a logger.
-
- .. code-block::
-
- Epoch 1: 4%|▎ | 40/1095 [00:03<01:37, 10.84it/s, loss=4.501, v_num=10]
-
- Here is an example how to override the defaults:
-
- .. code-block:: python
-
- def get_progress_bar_dict(self):
- # don't show the version number
- items = super().get_progress_bar_dict()
- items.pop("v_num", None)
- return items
-
- Return:
- Dictionary with the items to be displayed in the progress bar.
- """
- return progress_base.get_standard_metrics(self.trainer, self)
-
- def _verify_is_manual_optimization(self, fn_name):
- if self.automatic_optimization:
- raise MisconfigurationException(
- f"to use {fn_name}, please disable automatic optimization:"
- " set model property `automatic_optimization` as False"
- )
-
- @classmethod
- def _auto_collect_arguments(cls, frame=None) -> Tuple[Dict, Dict]:
- """Collect all module arguments in the current constructor and all child constructors. The child
- constructors are all the ``__init__`` methods that reach the current class through (chained)
- ``super().__init__()`` calls.
-
- Args:
- frame: instance frame
-
- Returns:
- self_arguments: arguments dictionary of the first instance
- parents_arguments: arguments dictionary of the parent's instances
- """
- if not frame:
- frame = inspect.currentframe()
-
- frame_args = collect_init_args(frame.f_back, [])
- self_arguments = frame_args[-1]
-
- # set hyper_parameters in child
- self_arguments = self_arguments
- parents_arguments = {}
-
- # add all arguments from parents
- for args in frame_args[:-1]:
- parents_arguments.update(args)
- return self_arguments, parents_arguments
-
-[docs] @torch.no_grad()
- def to_onnx(self, file_path: Union[str, Path], input_sample: Optional[Any] = None, **kwargs):
- """Saves the model in ONNX format.
-
- Args:
- file_path: The path of the file the onnx model should be saved to.
- input_sample: An input for tracing. Default: None (Use self.example_input_array)
- **kwargs: Will be passed to torch.onnx.export function.
-
- Example:
- >>> class SimpleModel(LightningModule):
- ... def __init__(self):
- ... super().__init__()
- ... self.l1 = torch.nn.Linear(in_features=64, out_features=4)
- ...
- ... def forward(self, x):
- ... return torch.relu(self.l1(x.view(x.size(0), -1)))
-
- >>> with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
- ... model = SimpleModel()
- ... input_sample = torch.randn((1, 64))
- ... model.to_onnx(tmpfile.name, input_sample, export_params=True)
- ... os.path.isfile(tmpfile.name)
- True
- """
- mode = self.training
-
- if input_sample is None:
- if self.example_input_array is None:
- raise ValueError(
- "Could not export to ONNX since neither `input_sample` nor"
- " `model.example_input_array` attribute is set."
- )
- input_sample = self.example_input_array
-
- input_sample = self._apply_batch_transfer_handler(input_sample)
-
- if not _TORCH_GREATER_EQUAL_1_10 and "example_outputs" not in kwargs:
- self.eval()
- if isinstance(input_sample, Tuple):
- kwargs["example_outputs"] = self(*input_sample)
- else:
- kwargs["example_outputs"] = self(input_sample)
-
- torch.onnx.export(self, input_sample, file_path, **kwargs)
- self.train(mode)
-
-[docs] @torch.no_grad()
- def to_torchscript(
- self,
- file_path: Optional[Union[str, Path]] = None,
- method: Optional[str] = "script",
- example_inputs: Optional[Any] = None,
- **kwargs,
- ) -> Union[ScriptModule, Dict[str, ScriptModule]]:
- """By default compiles the whole model to a :class:`~torch.jit.ScriptModule`. If you want to use tracing,
- please provided the argument ``method='trace'`` and make sure that either the `example_inputs` argument is
- provided, or the model has :attr:`example_input_array` set. If you would like to customize the modules that
- are scripted you should override this method. In case you want to return multiple modules, we recommend
- using a dictionary.
-
- Args:
- file_path: Path where to save the torchscript. Default: None (no file saved).
- method: Whether to use TorchScript's script or trace method. Default: 'script'
- example_inputs: An input to be used to do tracing when method is set to 'trace'.
- Default: None (uses :attr:`example_input_array`)
- **kwargs: Additional arguments that will be passed to the :func:`torch.jit.script` or
- :func:`torch.jit.trace` function.
-
- Note:
- - Requires the implementation of the
- :meth:`~pytorch_lightning.core.lightning.LightningModule.forward` method.
- - The exported script will be set to evaluation mode.
- - It is recommended that you install the latest supported version of PyTorch
- to use this feature without limitations. See also the :mod:`torch.jit`
- documentation for supported features.
-
- Example:
- >>> class SimpleModel(LightningModule):
- ... def __init__(self):
- ... super().__init__()
- ... self.l1 = torch.nn.Linear(in_features=64, out_features=4)
- ...
- ... def forward(self, x):
- ... return torch.relu(self.l1(x.view(x.size(0), -1)))
- ...
- >>> model = SimpleModel()
- >>> model.to_torchscript(file_path="model.pt") # doctest: +SKIP
- >>> os.path.isfile("model.pt") # doctest: +SKIP
- >>> torch.jit.save(model.to_torchscript(file_path="model_trace.pt", method='trace', # doctest: +SKIP
- ... example_inputs=torch.randn(1, 64))) # doctest: +SKIP
- >>> os.path.isfile("model_trace.pt") # doctest: +SKIP
- True
-
- Return:
- This LightningModule as a torchscript, regardless of whether `file_path` is
- defined or not.
- """
- mode = self.training
-
- self._running_torchscript = True
-
- if method == "script":
- torchscript_module = torch.jit.script(self.eval(), **kwargs)
- elif method == "trace":
- # if no example inputs are provided, try to see if model has example_input_array set
- if example_inputs is None:
- if self.example_input_array is None:
- raise ValueError(
- "Choosing method=`trace` requires either `example_inputs`"
- " or `model.example_input_array` to be defined."
- )
- example_inputs = self.example_input_array
-
- # automatically send example inputs to the right device and use trace
- example_inputs = self._apply_batch_transfer_handler(example_inputs)
- torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
- else:
- raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was: {method}")
-
- self.train(mode)
-
- if file_path is not None:
- fs = get_filesystem(file_path)
- with fs.open(file_path, "wb") as f:
- torch.jit.save(torchscript_module, f)
-
- self._running_torchscript = False
-
- return torchscript_module
-
- @property
- def model_size(self) -> float:
- """Returns the model size in MegaBytes (MB)
-
- Note:
- This property will not return correct value for Deepspeed (stage 3) and fully-sharded training.
- """
- if not self._running_torchscript: # remove with the deprecation removal
- rank_zero_deprecation(
- "The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7."
- " Please use the `pytorch_lightning.utilities.memory.get_model_size_mb`.",
- stacklevel=5,
- )
- return get_model_size_mb(self)
-
- @property
- def use_amp(self) -> bool:
- r"""
- .. deprecated:: v1.6.
-
- This property was deprecated in v1.6 and will be removed in v1.8.
- """
- if not self._running_torchscript: # remove with the deprecation removal
- rank_zero_deprecation(
- "`LightningModule.use_amp` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.amp_backend`.",
- stacklevel=5,
- )
- return self._use_amp
-
- @use_amp.setter
- def use_amp(self, use_amp: bool) -> None:
- r"""
- .. deprecated:: v1.6.
-
- This property was deprecated in v1.6 and will be removed in v1.8.
- """
- if not self._running_torchscript: # remove with the deprecation removal
- rank_zero_deprecation(
- "`LightningModule.use_amp` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.amp_backend`.",
- stacklevel=5,
- )
- self._use_amp = use_amp
-
-[docs] def add_to_queue(self, queue: pl.strategies.launchers.spawn._FakeQueue) -> None:
- """Appends the :attr:`trainer.callback_metrics` dictionary to the given queue. To avoid issues with memory
- sharing, we cast the data to numpy.
-
- Args:
- queue: the instance of the queue to append the data.
-
- .. deprecated:: v1.5
- This method was deprecated in v1.5 and will be removed in v1.7.
- """
-
-[docs] def get_from_queue(self, queue: pl.strategies.launchers.spawn._FakeQueue) -> None:
- """Retrieve the :attr:`trainer.callback_metrics` dictionary from the given queue. To preserve consistency,
- we cast back the data to ``torch.Tensor``.
-
- Args:
- queue: the instance of the queue from where to get the data.
-
- .. deprecated:: v1.5
- This method was deprecated in v1.5 and will be removed in v1.7.
- """
-
- @contextmanager
- def _prevent_trainer_and_dataloaders_deepcopy(self) -> None:
- self._should_prevent_trainer_and_dataloaders_deepcopy = True
- yield
- self._should_prevent_trainer_and_dataloaders_deepcopy = False
-
- def __getstate__(self) -> Dict[str, Any]:
- state = dict(self.__dict__)
- if self._should_prevent_trainer_and_dataloaders_deepcopy:
- state["trainer"] = None
- state.pop("train_dataloader", None)
- state.pop("val_dataloader", None)
- state.pop("test_dataloader", None)
- state.pop("predict_dataloader", None)
- return state
-
- def _register_sharded_tensor_state_dict_hooks_if_available(self) -> None:
- """Adds ShardedTensor state dict hooks if ShardedTensors are supported.
-
- These hooks ensure that ShardedTensors are included when saving, and are loaded the LightningModule correctly.
- """
- if not _TORCH_GREATER_EQUAL_1_10 or _IS_WINDOWS or not torch.distributed.is_available():
- rank_zero_debug("Could not register sharded tensor state dict hooks")
- return
-
- from torch.distributed._sharded_tensor import pre_load_state_dict_hook, state_dict_hook
-
- self._register_state_dict_hook(state_dict_hook)
- self._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Comet Logger
-------------
-"""
-
-import logging
-import os
-from argparse import Namespace
-from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union
-
-import torch
-from torch import is_tensor
-
-import pytorch_lightning as pl
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.imports import _module_available
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict
-from pytorch_lightning.utilities.rank_zero import rank_zero_only
-
-log = logging.getLogger(__name__)
-_COMET_AVAILABLE = _module_available("comet_ml")
-
-if _COMET_AVAILABLE:
- import comet_ml
- from comet_ml import ExistingExperiment as CometExistingExperiment
- from comet_ml import Experiment as CometExperiment
- from comet_ml import OfflineExperiment as CometOfflineExperiment
-
- try:
- from comet_ml.api import API
- except ModuleNotFoundError: # pragma: no-cover
- # For more information, see: https://www.comet.ml/docs/python-sdk/releases/#release-300
- from comet_ml.papi import API # pragma: no-cover
-else:
- # needed for test mocks, these tests shall be updated
- comet_ml = None
- CometExperiment, CometExistingExperiment, CometOfflineExperiment = None, None, None
- API = None
-
-
-[docs]class CometLogger(LightningLoggerBase):
- r"""
- Log using `Comet.ml <https://www.comet.ml>`_.
-
- Install it with pip:
-
- .. code-block:: bash
-
- pip install comet-ml
-
- Comet requires either an API Key (online mode) or a local directory path (offline mode).
-
- **ONLINE MODE**
-
- .. code-block:: python
-
- import os
- from pytorch_lightning import Trainer
- from pytorch_lightning.loggers import CometLogger
-
- # arguments made to CometLogger are passed on to the comet_ml.Experiment class
- comet_logger = CometLogger(
- api_key=os.environ.get("COMET_API_KEY"),
- workspace=os.environ.get("COMET_WORKSPACE"), # Optional
- save_dir=".", # Optional
- project_name="default_project", # Optional
- rest_api_key=os.environ.get("COMET_REST_API_KEY"), # Optional
- experiment_key=os.environ.get("COMET_EXPERIMENT_KEY"), # Optional
- experiment_name="lightning_logs", # Optional
- )
- trainer = Trainer(logger=comet_logger)
-
- **OFFLINE MODE**
-
- .. code-block:: python
-
- from pytorch_lightning.loggers import CometLogger
-
- # arguments made to CometLogger are passed on to the comet_ml.Experiment class
- comet_logger = CometLogger(
- save_dir=".",
- workspace=os.environ.get("COMET_WORKSPACE"), # Optional
- project_name="default_project", # Optional
- rest_api_key=os.environ.get("COMET_REST_API_KEY"), # Optional
- experiment_name="lightning_logs", # Optional
- )
- trainer = Trainer(logger=comet_logger)
-
- Args:
- api_key: Required in online mode. API key, found on Comet.ml. If not given, this
- will be loaded from the environment variable COMET_API_KEY or ~/.comet.config
- if either exists.
- save_dir: Required in offline mode. The path for the directory to save local
- comet logs. If given, this also sets the directory for saving checkpoints.
- project_name: Optional. Send your experiment to a specific project.
- Otherwise will be sent to Uncategorized Experiments.
- If the project name does not already exist, Comet.ml will create a new project.
- rest_api_key: Optional. Rest API key found in Comet.ml settings.
- This is used to determine version number
- experiment_name: Optional. String representing the name for this particular experiment on Comet.ml.
- experiment_key: Optional. If set, restores from existing experiment.
- offline: If api_key and save_dir are both given, this determines whether
- the experiment will be in online or offline mode. This is useful if you use
- save_dir to control the checkpoints directory and have a ~/.comet.config
- file but still want to run offline experiments.
- prefix: A string to put at the beginning of metric keys.
- \**kwargs: Additional arguments like `workspace`, `log_code`, etc. used by
- :class:`CometExperiment` can be passed as keyword arguments in this logger.
-
- Raises:
- ModuleNotFoundError:
- If required Comet package is not installed on the device.
- MisconfigurationException:
- If neither ``api_key`` nor ``save_dir`` are passed as arguments.
- """
-
- LOGGER_JOIN_CHAR = "-"
-
- def __init__(
- self,
- api_key: Optional[str] = None,
- save_dir: Optional[str] = None,
- project_name: Optional[str] = None,
- rest_api_key: Optional[str] = None,
- experiment_name: Optional[str] = None,
- experiment_key: Optional[str] = None,
- offline: bool = False,
- prefix: str = "",
- agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,
- agg_default_func: Optional[Callable[[Sequence[float]], float]] = None,
- **kwargs,
- ):
- if comet_ml is None:
- raise ModuleNotFoundError(
- "You want to use `comet_ml` logger which is not installed yet, install it with `pip install comet-ml`."
- )
- super().__init__(agg_key_funcs=agg_key_funcs, agg_default_func=agg_default_func)
- self._experiment = None
-
- # Determine online or offline mode based on which arguments were passed to CometLogger
- api_key = api_key or comet_ml.config.get_api_key(None, comet_ml.config.get_config())
-
- if api_key is not None and save_dir is not None:
- self.mode = "offline" if offline else "online"
- self.api_key = api_key
- self._save_dir = save_dir
- elif api_key is not None:
- self.mode = "online"
- self.api_key = api_key
- self._save_dir = None
- elif save_dir is not None:
- self.mode = "offline"
- self._save_dir = save_dir
- else:
- # If neither api_key nor save_dir are passed as arguments, raise an exception
- raise MisconfigurationException("CometLogger requires either api_key or save_dir during initialization.")
-
- log.info(f"CometLogger will be initialized in {self.mode} mode")
-
- self._project_name = project_name
- self._experiment_key = experiment_key
- self._experiment_name = experiment_name
- self._prefix = prefix
- self._kwargs = kwargs
- self._future_experiment_key = None
-
- if rest_api_key is not None:
- # Comet.ml rest API, used to determine version number
- self.rest_api_key = rest_api_key
- self.comet_api = API(self.rest_api_key)
- else:
- self.rest_api_key = None
- self.comet_api = None
-
- self._kwargs = kwargs
-
- @property
- @rank_zero_experiment
- def experiment(self) -> Union[CometExperiment, CometExistingExperiment, CometOfflineExperiment]:
- r"""
- Actual Comet object. To use Comet features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
-
- Example::
-
- self.logger.experiment.some_comet_function()
-
- """
- if self._experiment is not None:
- return self._experiment
-
- if self._future_experiment_key is not None:
- os.environ["COMET_EXPERIMENT_KEY"] = self._future_experiment_key
-
- try:
- if self.mode == "online":
- if self._experiment_key is None:
- self._experiment = CometExperiment(
- api_key=self.api_key, project_name=self._project_name, **self._kwargs
- )
- self._experiment_key = self._experiment.get_key()
- else:
- self._experiment = CometExistingExperiment(
- api_key=self.api_key,
- project_name=self._project_name,
- previous_experiment=self._experiment_key,
- **self._kwargs,
- )
- else:
- self._experiment = CometOfflineExperiment(
- offline_directory=self.save_dir, project_name=self._project_name, **self._kwargs
- )
- finally:
- if self._future_experiment_key is not None:
- os.environ.pop("COMET_EXPERIMENT_KEY")
- self._future_experiment_key = None
-
- if self._experiment_name:
- self._experiment.set_name(self._experiment_name)
-
- return self._experiment
-
-[docs] @rank_zero_only
- def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
- params = _convert_params(params)
- params = _flatten_dict(params)
- self.experiment.log_parameters(params)
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, Union[torch.Tensor, float]], step: Optional[int] = None) -> None:
- assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
- # Comet.ml expects metrics to be a dictionary of detached tensors on CPU
- metrics_without_epoch = metrics.copy()
- for key, val in metrics_without_epoch.items():
- if is_tensor(val):
- metrics_without_epoch[key] = val.cpu().detach()
-
- epoch = metrics_without_epoch.pop("epoch", None)
- metrics_without_epoch = _add_prefix(metrics_without_epoch, self._prefix, self.LOGGER_JOIN_CHAR)
- self.experiment.log_metrics(metrics_without_epoch, step=step, epoch=epoch)
-
- def reset_experiment(self):
- self._experiment = None
-
-[docs] @rank_zero_only
- def finalize(self, status: str) -> None:
- r"""
- When calling ``self.experiment.end()``, that experiment won't log any more data to Comet.
- That's why, if you need to log any more data, you need to create an ExistingCometExperiment.
- For example, to log data when testing your model after training, because when training is
- finalized :meth:`CometLogger.finalize` is called.
-
- This happens automatically in the :meth:`~CometLogger.experiment` property, when
- ``self._experiment`` is set to ``None``, i.e. ``self.reset_experiment()``.
- """
- self.experiment.end()
- self.reset_experiment()
-
- @property
- def save_dir(self) -> Optional[str]:
- """Gets the save directory.
-
- Returns:
- The path to the save directory.
- """
- return self._save_dir
-
- @property
- def name(self) -> str:
- """Gets the project name.
-
- Returns:
- The project name if it is specified, else "comet-default".
- """
- # Don't create an experiment if we don't have one
- if self._experiment is not None and self._experiment.project_name is not None:
- return self._experiment.project_name
-
- if self._project_name is not None:
- return self._project_name
-
- return "comet-default"
-
- @property
- def version(self) -> str:
- """Gets the version.
-
- Returns:
- The first one of the following that is set in the following order
-
- 1. experiment id.
- 2. experiment key.
- 3. "COMET_EXPERIMENT_KEY" environment variable.
- 4. future experiment key.
-
- If none are present generates a new guid.
- """
- # Don't create an experiment if we don't have one
- if self._experiment is not None:
- return self._experiment.id
-
- if self._experiment_key is not None:
- return self._experiment_key
-
- if "COMET_EXPERIMENT_KEY" in os.environ:
- return os.environ["COMET_EXPERIMENT_KEY"]
-
- if self._future_experiment_key is not None:
- return self._future_experiment_key
-
- # Pre-generate an experiment key
- self._future_experiment_key = comet_ml.generate_guid()
-
- return self._future_experiment_key
-
- def __getstate__(self):
- state = self.__dict__.copy()
-
- # Save the experiment id in case an experiment object already exists,
- # this way we could create an ExistingExperiment pointing to the same
- # experiment
- state["_experiment_key"] = self._experiment.id if self._experiment is not None else None
-
- # Remove the experiment object as it contains hard to pickle objects
- # (like network connections), the experiment object will be recreated if
- # needed later
- state["_experiment"] = None
- return state
-
-[docs] def log_graph(self, model: "pl.LightningModule", input_array=None) -> None:
- if self._experiment is not None:
- self._experiment.set_model_graph(model)
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-CSV logger
-----------
-
-CSV logger for basic experiment logging that does not require opening ports
-
-"""
-import csv
-import logging
-import os
-from argparse import Namespace
-from typing import Any, Dict, Optional, Union
-
-import torch
-
-from pytorch_lightning.core.saving import save_hparams_to_yaml
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params
-from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
-
-log = logging.getLogger(__name__)
-
-
-class ExperimentWriter:
- r"""
- Experiment writer for CSVLogger.
-
- Currently supports to log hyperparameters and metrics in YAML and CSV
- format, respectively.
-
- Args:
- log_dir: Directory for the experiment logs
- """
-
- NAME_HPARAMS_FILE = "hparams.yaml"
- NAME_METRICS_FILE = "metrics.csv"
-
- def __init__(self, log_dir: str) -> None:
- self.hparams = {}
- self.metrics = []
-
- self.log_dir = log_dir
- if os.path.exists(self.log_dir) and os.listdir(self.log_dir):
- rank_zero_warn(
- f"Experiment logs directory {self.log_dir} exists and is not empty."
- " Previous log files in this directory will be deleted when the new ones are saved!"
- )
- os.makedirs(self.log_dir, exist_ok=True)
-
- self.metrics_file_path = os.path.join(self.log_dir, self.NAME_METRICS_FILE)
-
- def log_hparams(self, params: Dict[str, Any]) -> None:
- """Record hparams."""
- self.hparams.update(params)
-
- def log_metrics(self, metrics_dict: Dict[str, float], step: Optional[int] = None) -> None:
- """Record metrics."""
-
- def _handle_value(value):
- if isinstance(value, torch.Tensor):
- return value.item()
- return value
-
- if step is None:
- step = len(self.metrics)
-
- metrics = {k: _handle_value(v) for k, v in metrics_dict.items()}
- metrics["step"] = step
- self.metrics.append(metrics)
-
- def save(self) -> None:
- """Save recorded hparams and metrics into files."""
- hparams_file = os.path.join(self.log_dir, self.NAME_HPARAMS_FILE)
- save_hparams_to_yaml(hparams_file, self.hparams)
-
- if not self.metrics:
- return
-
- last_m = {}
- for m in self.metrics:
- last_m.update(m)
- metrics_keys = list(last_m.keys())
-
- with open(self.metrics_file_path, "w", newline="") as f:
- writer = csv.DictWriter(f, fieldnames=metrics_keys)
- writer.writeheader()
- writer.writerows(self.metrics)
-
-
-[docs]class CSVLogger(LightningLoggerBase):
- r"""
- Log to local file system in yaml and CSV format.
-
- Logs are saved to ``os.path.join(save_dir, name, version)``.
-
- Example:
- >>> from pytorch_lightning import Trainer
- >>> from pytorch_lightning.loggers import CSVLogger
- >>> logger = CSVLogger("logs", name="my_exp_name")
- >>> trainer = Trainer(logger=logger)
-
- Args:
- save_dir: Save directory
- name: Experiment name. Defaults to ``'default'``.
- version: Experiment version. If version is not specified the logger inspects the save
- directory for existing versions, then automatically assigns the next available version.
- prefix: A string to put at the beginning of metric keys.
- flush_logs_every_n_steps: How often to flush logs to disk (defaults to every 100 steps).
- """
-
- LOGGER_JOIN_CHAR = "-"
-
- def __init__(
- self,
- save_dir: str,
- name: Optional[str] = "lightning_logs",
- version: Optional[Union[int, str]] = None,
- prefix: str = "",
- flush_logs_every_n_steps: int = 100,
- ):
- super().__init__()
- self._save_dir = save_dir
- self._name = name or ""
- self._version = version
- self._prefix = prefix
- self._experiment = None
- self._flush_logs_every_n_steps = flush_logs_every_n_steps
-
- @property
- def root_dir(self) -> str:
- """Parent directory for all checkpoint subdirectories.
-
- If the experiment name parameter is an empty string, no experiment subdirectory is used and the checkpoint will
- be saved in "save_dir/version"
- """
- return os.path.join(self.save_dir, self.name)
-
- @property
- def log_dir(self) -> str:
- """The log directory for this run.
-
- By default, it is named ``'version_${self.version}'`` but it can be overridden by passing a string value for the
- constructor's version parameter instead of ``None`` or an int.
- """
- # create a pseudo standard path
- version = self.version if isinstance(self.version, str) else f"version_{self.version}"
- log_dir = os.path.join(self.root_dir, version)
- return log_dir
-
- @property
- def save_dir(self) -> Optional[str]:
- """The current directory where logs are saved.
-
- Returns:
- The path to current directory where logs are saved.
- """
- return self._save_dir
-
- @property
- @rank_zero_experiment
- def experiment(self) -> ExperimentWriter:
- r"""
-
- Actual ExperimentWriter object. To use ExperimentWriter features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
-
- Example::
-
- self.logger.experiment.some_experiment_writer_function()
-
- """
- if self._experiment:
- return self._experiment
-
- os.makedirs(self.root_dir, exist_ok=True)
- self._experiment = ExperimentWriter(log_dir=self.log_dir)
- return self._experiment
-
-[docs] @rank_zero_only
- def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
- params = _convert_params(params)
- self.experiment.log_hparams(params)
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
- metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
- self.experiment.log_metrics(metrics, step)
- if step is not None and (step + 1) % self._flush_logs_every_n_steps == 0:
- self.save()
-
-
-
-
-
- @property
- def name(self) -> str:
- """Gets the name of the experiment.
-
- Returns:
- The name of the experiment.
- """
- return self._name
-
- @property
- def version(self) -> int:
- """Gets the version of the experiment.
-
- Returns:
- The version of the experiment if it is specified, else the next version.
- """
- if self._version is None:
- self._version = self._get_next_version()
- return self._version
-
- def _get_next_version(self):
- root_dir = self.root_dir
-
- if not os.path.isdir(root_dir):
- log.warning("Missing logger folder: %s", root_dir)
- return 0
-
- existing_versions = []
- for d in os.listdir(root_dir):
- if os.path.isdir(os.path.join(root_dir, d)) and d.startswith("version_"):
- existing_versions.append(int(d.split("_")[1]))
-
- if len(existing_versions) == 0:
- return 0
-
- return max(existing_versions) + 1
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-MLflow Logger
--------------
-"""
-import logging
-import os
-import re
-from argparse import Namespace
-from time import time
-from typing import Any, Dict, Optional, Union
-
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.imports import _module_available
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict
-from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
-
-log = logging.getLogger(__name__)
-LOCAL_FILE_URI_PREFIX = "file:"
-_MLFLOW_AVAILABLE = _module_available("mlflow")
-try:
- import mlflow
- from mlflow.tracking import context, MlflowClient
- from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
-# todo: there seems to be still some remaining import error with Conda env
-except ModuleNotFoundError:
- _MLFLOW_AVAILABLE = False
- mlflow, MlflowClient, context = None, None, None
- MLFLOW_RUN_NAME = "mlflow.runName"
-
-# before v1.1.0
-if hasattr(context, "resolve_tags"):
- from mlflow.tracking.context import resolve_tags
-
-
-# since v1.1.0
-elif hasattr(context, "registry"):
- from mlflow.tracking.context.registry import resolve_tags
-else:
-
- def resolve_tags(tags=None):
- return tags
-
-
-[docs]class MLFlowLogger(LightningLoggerBase):
- """Log using `MLflow <https://mlflow.org>`_.
-
- Install it with pip:
-
- .. code-block:: bash
-
- pip install mlflow
-
- .. code-block:: python
-
- from pytorch_lightning import Trainer
- from pytorch_lightning.loggers import MLFlowLogger
-
- mlf_logger = MLFlowLogger(experiment_name="lightning_logs", tracking_uri="file:./ml-runs")
- trainer = Trainer(logger=mlf_logger)
-
- Use the logger anywhere in your :class:`~pytorch_lightning.core.lightning.LightningModule` as follows:
-
- .. code-block:: python
-
- from pytorch_lightning import LightningModule
-
-
- class LitModel(LightningModule):
- def training_step(self, batch, batch_idx):
- # example
- self.logger.experiment.whatever_ml_flow_supports(...)
-
- def any_lightning_module_function_or_hook(self):
- self.logger.experiment.whatever_ml_flow_supports(...)
-
- Args:
- experiment_name: The name of the experiment.
- run_name: Name of the new run. The `run_name` is internally stored as a ``mlflow.runName`` tag.
- If the ``mlflow.runName`` tag has already been set in `tags`, the value is overridden by the `run_name`.
- tracking_uri: Address of local or remote tracking server.
- If not provided, defaults to `MLFLOW_TRACKING_URI` environment variable if set, otherwise it falls
- back to `file:<save_dir>`.
- tags: A dictionary tags for the experiment.
- save_dir: A path to a local directory where the MLflow runs get saved.
- Defaults to `./mlflow` if `tracking_uri` is not provided.
- Has no effect if `tracking_uri` is provided.
- prefix: A string to put at the beginning of metric keys.
- artifact_location: The location to store run artifacts. If not provided, the server picks an appropriate
- default.
- run_id: The run identifier of the experiment. If not provided, a new run is started.
-
- Raises:
- ModuleNotFoundError:
- If required MLFlow package is not installed on the device.
- """
-
- LOGGER_JOIN_CHAR = "-"
-
- def __init__(
- self,
- experiment_name: str = "lightning_logs",
- run_name: Optional[str] = None,
- tracking_uri: Optional[str] = os.getenv("MLFLOW_TRACKING_URI"),
- tags: Optional[Dict[str, Any]] = None,
- save_dir: Optional[str] = "./mlruns",
- prefix: str = "",
- artifact_location: Optional[str] = None,
- run_id: Optional[str] = None,
- ):
- if mlflow is None:
- raise ModuleNotFoundError(
- "You want to use `mlflow` logger which is not installed yet, install it with `pip install mlflow`."
- )
- super().__init__()
- if not tracking_uri:
- tracking_uri = f"{LOCAL_FILE_URI_PREFIX}{save_dir}"
-
- self._experiment_name = experiment_name
- self._experiment_id = None
- self._tracking_uri = tracking_uri
- self._run_name = run_name
- self._run_id = run_id
- self.tags = tags
- self._prefix = prefix
- self._artifact_location = artifact_location
-
- self._initialized = False
-
- self._mlflow_client = MlflowClient(tracking_uri)
-
- @property
- @rank_zero_experiment
- def experiment(self) -> MlflowClient:
- r"""
- Actual MLflow object. To use MLflow features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
-
- Example::
-
- self.logger.experiment.some_mlflow_function()
-
- """
-
- if self._initialized:
- return self._mlflow_client
-
- if self._run_id is not None:
- run = self._mlflow_client.get_run(self._run_id)
- self._experiment_id = run.info.experiment_id
- self._initialized = True
- return self._mlflow_client
-
- if self._experiment_id is None:
- expt = self._mlflow_client.get_experiment_by_name(self._experiment_name)
- if expt is not None:
- self._experiment_id = expt.experiment_id
- else:
- log.warning(f"Experiment with name {self._experiment_name} not found. Creating it.")
- self._experiment_id = self._mlflow_client.create_experiment(
- name=self._experiment_name, artifact_location=self._artifact_location
- )
-
- if self._run_id is None:
- if self._run_name is not None:
- self.tags = self.tags or {}
- if MLFLOW_RUN_NAME in self.tags:
- log.warning(
- f"The tag {MLFLOW_RUN_NAME} is found in tags. The value will be overridden by {self._run_name}."
- )
- self.tags[MLFLOW_RUN_NAME] = self._run_name
- run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=resolve_tags(self.tags))
- self._run_id = run.info.run_id
- self._initialized = True
- return self._mlflow_client
-
- @property
- def run_id(self) -> str:
- """Create the experiment if it does not exist to get the run id.
-
- Returns:
- The run id.
- """
- _ = self.experiment
- return self._run_id
-
- @property
- def experiment_id(self) -> str:
- """Create the experiment if it does not exist to get the experiment id.
-
- Returns:
- The experiment id.
- """
- _ = self.experiment
- return self._experiment_id
-
-[docs] @rank_zero_only
- def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
- params = _convert_params(params)
- params = _flatten_dict(params)
- for k, v in params.items():
- if len(str(v)) > 250:
- rank_zero_warn(
- f"Mlflow only allows parameters with up to 250 characters. Discard {k}={v}", category=RuntimeWarning
- )
- continue
-
- self.experiment.log_param(self.run_id, k, v)
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
- assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
-
- metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
-
- timestamp_ms = int(time() * 1000)
- for k, v in metrics.items():
- if isinstance(v, str):
- log.warning(f"Discarding metric with string value {k}={v}.")
- continue
-
- new_k = re.sub("[^a-zA-Z0-9_/. -]+", "", k)
- if k != new_k:
- rank_zero_warn(
- "MLFlow only allows '_', '/', '.' and ' ' special characters in metric name."
- f" Replacing {k} with {new_k}.",
- category=RuntimeWarning,
- )
- k = new_k
-
- self.experiment.log_metric(self.run_id, k, v, timestamp_ms, step)
-
-[docs] @rank_zero_only
- def finalize(self, status: str = "FINISHED") -> None:
- super().finalize(status)
- status = "FINISHED" if status == "success" else status
- if self.experiment.get_run(self.run_id):
- self.experiment.set_terminated(self.run_id, status)
-
- @property
- def save_dir(self) -> Optional[str]:
- """The root file directory in which MLflow experiments are saved.
-
- Return:
- Local path to the root experiment directory if the tracking uri is local.
- Otherwise returns `None`.
- """
- if self._tracking_uri.startswith(LOCAL_FILE_URI_PREFIX):
- return self._tracking_uri.lstrip(LOCAL_FILE_URI_PREFIX)
-
- @property
- def name(self) -> str:
- """Get the experiment id.
-
- Returns:
- The experiment id.
- """
- return self.experiment_id
-
- @property
- def version(self) -> str:
- """Get the run id.
-
- Returns:
- The run id.
- """
- return self.run_id
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Neptune Logger
---------------
-"""
-__all__ = [
- "NeptuneLogger",
-]
-
-import logging
-import os
-import warnings
-from argparse import Namespace
-from functools import reduce
-from typing import Any, Callable, Dict, Generator, Mapping, Optional, Sequence, Set, Union
-from weakref import ReferenceType
-
-import torch
-
-from pytorch_lightning import __version__
-from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.imports import _NEPTUNE_AVAILABLE, _NEPTUNE_GREATER_EQUAL_0_9
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _sanitize_callable_params
-from pytorch_lightning.utilities.model_summary import ModelSummary
-from pytorch_lightning.utilities.rank_zero import rank_zero_only
-
-if _NEPTUNE_AVAILABLE and _NEPTUNE_GREATER_EQUAL_0_9:
- try:
- from neptune import new as neptune
- from neptune.new.exceptions import NeptuneLegacyProjectException, NeptuneOfflineModeFetchException
- from neptune.new.run import Run
- from neptune.new.types import File as NeptuneFile
- except ModuleNotFoundError:
- import neptune
- from neptune.exceptions import NeptuneLegacyProjectException, NeptuneOfflineModeFetchException
- from neptune.run import Run
- from neptune.types import File as NeptuneFile
-else:
- # needed for test mocks, and function signatures
- neptune, Run, NeptuneFile = None, None, None
-
-log = logging.getLogger(__name__)
-
-_INTEGRATION_VERSION_KEY = "source_code/integrations/pytorch-lightning"
-
-# kwargs used in previous NeptuneLogger version, now deprecated
-_LEGACY_NEPTUNE_INIT_KWARGS = [
- "project_name",
- "offline_mode",
- "experiment_name",
- "experiment_id",
- "params",
- "properties",
- "upload_source_files",
- "abort_callback",
- "logger",
- "upload_stdout",
- "upload_stderr",
- "send_hardware_metrics",
- "run_monitoring_thread",
- "handle_uncaught_exceptions",
- "git_info",
- "hostname",
- "notebook_id",
- "notebook_path",
-]
-
-# kwargs used in legacy NeptuneLogger from neptune-pytorch-lightning package
-_LEGACY_NEPTUNE_LOGGER_KWARGS = [
- "base_namespace",
- "close_after_fit",
-]
-
-
-[docs]class NeptuneLogger(LightningLoggerBase):
- r"""
- Log using `Neptune <https://neptune.ai>`_.
-
- Install it with pip:
-
- .. code-block:: bash
-
- pip install neptune-client
-
- or conda:
-
- .. code-block:: bash
-
- conda install -c conda-forge neptune-client
-
- **Quickstart**
-
- Pass NeptuneLogger instance to the Trainer to log metadata with Neptune:
-
- .. code-block:: python
-
-
- from pytorch_lightning import Trainer
- from pytorch_lightning.loggers import NeptuneLogger
-
- neptune_logger = NeptuneLogger(
- api_key="ANONYMOUS", # replace with your own
- project="common/pytorch-lightning-integration", # format "<WORKSPACE/PROJECT>"
- tags=["training", "resnet"], # optional
- )
- trainer = Trainer(max_epochs=10, logger=neptune_logger)
-
- **How to use NeptuneLogger?**
-
- Use the logger anywhere in your :class:`~pytorch_lightning.core.lightning.LightningModule` as follows:
-
- .. code-block:: python
-
- from neptune.new.types import File
- from pytorch_lightning import LightningModule
-
-
- class LitModel(LightningModule):
- def training_step(self, batch, batch_idx):
- # log metrics
- acc = ...
- self.log("train/loss", loss)
-
- def any_lightning_module_function_or_hook(self):
- # log images
- img = ...
- self.logger.experiment["train/misclassified_images"].log(File.as_image(img))
-
- # generic recipe
- metadata = ...
- self.logger.experiment["your/metadata/structure"].log(metadata)
-
- Note that syntax: ``self.logger.experiment["your/metadata/structure"].log(metadata)`` is specific to Neptune
- and it extends logger capabilities. Specifically, it allows you to log various types of metadata
- like scores, files, images, interactive visuals, CSVs, etc.
- Refer to the `Neptune docs <https://docs.neptune.ai/you-should-know/logging-metadata#essential-logging-methods>`_
- for more detailed explanations.
- You can also use regular logger methods ``log_metrics()``, and ``log_hyperparams()`` with NeptuneLogger
- as these are also supported.
-
- **Log after fitting or testing is finished**
-
- You can log objects after the fitting or testing methods are finished:
-
- .. code-block:: python
-
- neptune_logger = NeptuneLogger(project="common/pytorch-lightning-integration")
-
- trainer = pl.Trainer(logger=neptune_logger)
- model = ...
- datamodule = ...
- trainer.fit(model, datamodule=datamodule)
- trainer.test(model, datamodule=datamodule)
-
- # Log objects after `fit` or `test` methods
- # model summary
- neptune_logger.log_model_summary(model=model, max_depth=-1)
-
- # generic recipe
- metadata = ...
- neptune_logger.experiment["your/metadata/structure"].log(metadata)
-
- **Log model checkpoints**
-
- If you have :class:`~pytorch_lightning.callbacks.ModelCheckpoint` configured,
- Neptune logger automatically logs model checkpoints.
- Model weights will be uploaded to the: "model/checkpoints" namespace in the Neptune Run.
- You can disable this option:
-
- .. code-block:: python
-
- neptune_logger = NeptuneLogger(project="common/pytorch-lightning-integration", log_model_checkpoints=False)
-
- **Pass additional parameters to the Neptune run**
-
- You can also pass ``neptune_run_kwargs`` to specify the run in the greater detail, like ``tags`` or ``description``:
-
- .. testcode::
-
- from pytorch_lightning import Trainer
- from pytorch_lightning.loggers import NeptuneLogger
-
- neptune_logger = NeptuneLogger(
- project="common/pytorch-lightning-integration",
- name="lightning-run",
- description="mlp quick run with pytorch-lightning",
- tags=["mlp", "quick-run"],
- )
- trainer = Trainer(max_epochs=3, logger=neptune_logger)
-
- Check `run documentation <https://docs.neptune.ai/essentials/api-reference/run>`_
- for more info about additional run parameters.
-
- **Details about Neptune run structure**
-
- Runs can be viewed as nested dictionary-like structures that you can define in your code.
- Thanks to this you can easily organize your metadata in a way that is most convenient for you.
-
- The hierarchical structure that you apply to your metadata will be reflected later in the UI.
-
- You can organize this way any type of metadata - images, parameters, metrics, model checkpoint, CSV files, etc.
-
- See Also:
- - Read about
- `what object you can log to Neptune <https://docs.neptune.ai/you-should-know/what-can-you-log-and-display>`_.
- - Check `example run <https://app.neptune.ai/o/common/org/pytorch-lightning-integration/e/PTL-1/all>`_
- with multiple types of metadata logged.
- - For more detailed info check
- `user guide <https://docs.neptune.ai/integrations-and-supported-tools/model-training/pytorch-lightning>`_.
-
- Args:
- api_key: Optional.
- Neptune API token, found on https://neptune.ai upon registration.
- Read: `how to find and set Neptune API token <https://docs.neptune.ai/administration/security-and-privacy/
- how-to-find-and-set-neptune-api-token>`_.
- It is recommended to keep it in the `NEPTUNE_API_TOKEN`
- environment variable and then you can drop ``api_key=None``.
- project: Optional.
- Name of a project in a form of "my_workspace/my_project" for example "tom/mask-rcnn".
- If ``None``, the value of `NEPTUNE_PROJECT` environment variable will be taken.
- You need to create the project in https://neptune.ai first.
- name: Optional. Editable name of the run.
- Run name appears in the "all metadata/sys" section in Neptune UI.
- run: Optional. Default is ``None``. The Neptune ``Run`` object.
- If specified, this `Run`` will be used for logging, instead of a new Run.
- When run object is passed you can't specify other neptune properties.
- log_model_checkpoints: Optional. Default is ``True``. Log model checkpoint to Neptune.
- Works only if ``ModelCheckpoint`` is passed to the ``Trainer``.
- prefix: Optional. Default is ``"training"``. Root namespace for all metadata logging.
- \**neptune_run_kwargs: Additional arguments like ``tags``, ``description``, ``capture_stdout``, etc.
- used when run is created.
-
- Raises:
- ModuleNotFoundError:
- If required Neptune package in version >=0.9 is not installed on the device.
- TypeError:
- If configured project has not been migrated to new structure yet.
- ValueError:
- If argument passed to the logger's constructor is incorrect.
- """
-
- LOGGER_JOIN_CHAR = "/"
- PARAMETERS_KEY = "hyperparams"
- ARTIFACTS_KEY = "artifacts"
-
- def __init__(
- self,
- *, # force users to call `NeptuneLogger` initializer with `kwargs`
- api_key: Optional[str] = None,
- project: Optional[str] = None,
- name: Optional[str] = None,
- run: Optional["Run"] = None,
- log_model_checkpoints: Optional[bool] = True,
- prefix: str = "training",
- agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,
- agg_default_func: Optional[Callable[[Sequence[float]], float]] = None,
- **neptune_run_kwargs,
- ):
- # verify if user passed proper init arguments
- self._verify_input_arguments(api_key, project, name, run, neptune_run_kwargs)
- if neptune is None:
- raise ModuleNotFoundError(
- "You want to use the `Neptune` logger which is not installed yet, install it with"
- " `pip install neptune-client`."
- )
-
- super().__init__(agg_key_funcs=agg_key_funcs, agg_default_func=agg_default_func)
- self._log_model_checkpoints = log_model_checkpoints
- self._prefix = prefix
- self._run_name = name
- self._project_name = project
- self._api_key = api_key
- self._run_instance = run
- self._neptune_run_kwargs = neptune_run_kwargs
- self._run_short_id = None
-
- if self._run_instance is not None:
- self._retrieve_run_data()
-
- # make sure that we've log integration version for outside `Run` instances
- self._run_instance[_INTEGRATION_VERSION_KEY] = __version__
-
- def _retrieve_run_data(self):
- try:
- self._run_instance.wait()
- self._run_short_id = self._run_instance["sys/id"].fetch()
- self._run_name = self._run_instance["sys/name"].fetch()
- except NeptuneOfflineModeFetchException:
- self._run_short_id = "OFFLINE"
- self._run_name = "offline-name"
-
- @property
- def _neptune_init_args(self):
- args = {}
- # Backward compatibility in case of previous version retrieval
- try:
- args = self._neptune_run_kwargs
- except AttributeError:
- pass
-
- if self._project_name is not None:
- args["project"] = self._project_name
-
- if self._api_key is not None:
- args["api_token"] = self._api_key
-
- if self._run_short_id is not None:
- args["run"] = self._run_short_id
-
- # Backward compatibility in case of previous version retrieval
- try:
- if self._run_name is not None:
- args["name"] = self._run_name
- except AttributeError:
- pass
-
- return args
-
- def _construct_path_with_prefix(self, *keys) -> str:
- """Return sequence of keys joined by `LOGGER_JOIN_CHAR`, started with `_prefix` if defined."""
- if self._prefix:
- return self.LOGGER_JOIN_CHAR.join([self._prefix, *keys])
- return self.LOGGER_JOIN_CHAR.join(keys)
-
- @staticmethod
- def _verify_input_arguments(
- api_key: Optional[str],
- project: Optional[str],
- name: Optional[str],
- run: Optional["Run"],
- neptune_run_kwargs: dict,
- ):
- legacy_kwargs_msg = (
- "Following kwargs are deprecated: {legacy_kwargs}.\n"
- "If you are looking for the Neptune logger using legacy Python API,"
- " it's still available as part of neptune-contrib package:\n"
- " - https://docs-legacy.neptune.ai/integrations/pytorch_lightning.html\n"
- "The NeptuneLogger was re-written to use the neptune.new Python API\n"
- " - https://neptune.ai/blog/neptune-new\n"
- " - https://docs.neptune.ai/integrations-and-supported-tools/model-training/pytorch-lightning\n"
- "You should use arguments accepted by either NeptuneLogger.init() or neptune.init()"
- )
-
- # check if user used legacy kwargs expected in `NeptuneLegacyLogger`
- used_legacy_kwargs = [
- legacy_kwarg for legacy_kwarg in neptune_run_kwargs if legacy_kwarg in _LEGACY_NEPTUNE_INIT_KWARGS
- ]
- if used_legacy_kwargs:
- raise ValueError(legacy_kwargs_msg.format(legacy_kwargs=used_legacy_kwargs))
-
- # check if user used legacy kwargs expected in `NeptuneLogger` from neptune-pytorch-lightning package
- used_legacy_neptune_kwargs = [
- legacy_kwarg for legacy_kwarg in neptune_run_kwargs if legacy_kwarg in _LEGACY_NEPTUNE_LOGGER_KWARGS
- ]
- if used_legacy_neptune_kwargs:
- raise ValueError(legacy_kwargs_msg.format(legacy_kwargs=used_legacy_neptune_kwargs))
-
- # check if user passed new client `Run` object
- if run is not None and not isinstance(run, Run):
- raise ValueError(
- "Run parameter expected to be of type `neptune.new.Run`.\n"
- "If you are looking for the Neptune logger using legacy Python API,"
- " it's still available as part of neptune-contrib package:\n"
- " - https://docs-legacy.neptune.ai/integrations/pytorch_lightning.html\n"
- "The NeptuneLogger was re-written to use the neptune.new Python API\n"
- " - https://neptune.ai/blog/neptune-new\n"
- " - https://docs.neptune.ai/integrations-and-supported-tools/model-training/pytorch-lightning\n"
- )
-
- # check if user passed redundant neptune.init arguments when passed run
- any_neptune_init_arg_passed = any(arg is not None for arg in [api_key, project, name]) or neptune_run_kwargs
- if run is not None and any_neptune_init_arg_passed:
- raise ValueError(
- "When an already initialized run object is provided"
- " you can't provide other neptune.init() parameters.\n"
- )
-
- def __getstate__(self):
- state = self.__dict__.copy()
- # Run instance can't be pickled
- state["_run_instance"] = None
- return state
-
- def __setstate__(self, state):
- self.__dict__ = state
- self._run_instance = neptune.init(**self._neptune_init_args)
-
- @property
- @rank_zero_experiment
- def experiment(self) -> Run:
- r"""
- Actual Neptune run object. Allows you to use neptune logging features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule`.
-
- Example::
-
- class LitModel(LightningModule):
- def training_step(self, batch, batch_idx):
- # log metrics
- acc = ...
- self.logger.experiment["train/acc"].log(acc)
-
- # log images
- img = ...
- self.logger.experiment["train/misclassified_images"].log(File.as_image(img))
-
- Note that syntax: ``self.logger.experiment["your/metadata/structure"].log(metadata)``
- is specific to Neptune and it extends logger capabilities.
- Specifically, it allows you to log various types of metadata like scores, files,
- images, interactive visuals, CSVs, etc. Refer to the
- `Neptune docs <https://docs.neptune.ai/you-should-know/logging-metadata#essential-logging-methods>`_
- for more detailed explanations.
- You can also use regular logger methods ``log_metrics()``, and ``log_hyperparams()``
- with NeptuneLogger as these are also supported.
- """
- return self.run
-
- @property
- @rank_zero_experiment
- def run(self) -> Run:
- try:
- if not self._run_instance:
- self._run_instance = neptune.init(**self._neptune_init_args)
- self._retrieve_run_data()
- # make sure that we've log integration version for newly created
- self._run_instance[_INTEGRATION_VERSION_KEY] = __version__
-
- return self._run_instance
- except NeptuneLegacyProjectException as e:
- raise TypeError(
- f"Project {self._project_name} has not been migrated to the new structure."
- " You can still integrate it with the Neptune logger using legacy Python API"
- " available as part of neptune-contrib package:"
- " https://docs-legacy.neptune.ai/integrations/pytorch_lightning.html\n"
- ) from e
-
-[docs] @rank_zero_only
- def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None: # skipcq: PYL-W0221
- r"""
- Log hyper-parameters to the run.
-
- Hyperparams will be logged under the "<prefix>/hyperparams" namespace.
-
- Note:
-
- You can also log parameters by directly using the logger instance:
- ``neptune_logger.experiment["model/hyper-parameters"] = params_dict``.
-
- In this way you can keep hierarchical structure of the parameters.
-
- Args:
- params: `dict`.
- Python dictionary structure with parameters.
-
- Example::
-
- from pytorch_lightning.loggers import NeptuneLogger
-
- PARAMS = {
- "batch_size": 64,
- "lr": 0.07,
- "decay_factor": 0.97
- }
-
- neptune_logger = NeptuneLogger(
- api_key="ANONYMOUS",
- project="common/pytorch-lightning-integration"
- )
-
- neptune_logger.log_hyperparams(PARAMS)
- """
- params = _convert_params(params)
- params = _sanitize_callable_params(params)
-
- parameters_key = self.PARAMETERS_KEY
- parameters_key = self._construct_path_with_prefix(parameters_key)
-
- self.run[parameters_key] = params
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, Union[torch.Tensor, float]], step: Optional[int] = None) -> None:
- """Log metrics (numeric values) in Neptune runs.
-
- Args:
- metrics: Dictionary with metric names as keys and measured quantities as values.
- step: Step number at which the metrics should be recorded, currently ignored.
- """
- if rank_zero_only.rank != 0:
- raise ValueError("run tried to log from global_rank != 0")
-
- metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
-
- for key, val in metrics.items():
- # `step` is ignored because Neptune expects strictly increasing step values which
- # Lightning does not always guarantee.
- self.run[key].log(val)
-
-[docs] @rank_zero_only
- def finalize(self, status: str) -> None:
- if status:
- self.run[self._construct_path_with_prefix("status")] = status
-
- super().finalize(status)
-
- @property
- def save_dir(self) -> Optional[str]:
- """Gets the save directory of the experiment which in this case is ``None`` because Neptune does not save
- locally.
-
- Returns:
- the root directory where experiment logs get saved
- """
- return os.path.join(os.getcwd(), ".neptune")
-
- @rank_zero_only
- def log_model_summary(self, model, max_depth=-1):
- model_str = str(ModelSummary(model=model, max_depth=max_depth))
- self.run[self._construct_path_with_prefix("model/summary")] = neptune.types.File.from_content(
- content=model_str, extension="txt"
- )
-
-[docs] @rank_zero_only
- def after_save_checkpoint(self, checkpoint_callback: "ReferenceType[ModelCheckpoint]") -> None:
- """Automatically log checkpointed model. Called after model checkpoint callback saves a new checkpoint.
-
- Args:
- checkpoint_callback: the model checkpoint callback instance
- """
- if not self._log_model_checkpoints:
- return
-
- file_names = set()
- checkpoints_namespace = self._construct_path_with_prefix("model/checkpoints")
-
- # save last model
- if checkpoint_callback.last_model_path:
- model_last_name = self._get_full_model_name(checkpoint_callback.last_model_path, checkpoint_callback)
- file_names.add(model_last_name)
- self.run[f"{checkpoints_namespace}/{model_last_name}"].upload(checkpoint_callback.last_model_path)
-
- # save best k models
- for key in checkpoint_callback.best_k_models.keys():
- model_name = self._get_full_model_name(key, checkpoint_callback)
- file_names.add(model_name)
- self.run[f"{checkpoints_namespace}/{model_name}"].upload(key)
-
- # log best model path and checkpoint
- if checkpoint_callback.best_model_path:
- self.run[self._construct_path_with_prefix("model/best_model_path")] = checkpoint_callback.best_model_path
-
- model_name = self._get_full_model_name(checkpoint_callback.best_model_path, checkpoint_callback)
- file_names.add(model_name)
- self.run[f"{checkpoints_namespace}/{model_name}"].upload(checkpoint_callback.best_model_path)
-
- # remove old models logged to experiment if they are not part of best k models at this point
- if self.run.exists(checkpoints_namespace):
- exp_structure = self.run.get_structure()
- uploaded_model_names = self._get_full_model_names_from_exp_structure(exp_structure, checkpoints_namespace)
-
- for file_to_drop in list(uploaded_model_names - file_names):
- del self.run[f"{checkpoints_namespace}/{file_to_drop}"]
-
- # log best model score
- if checkpoint_callback.best_model_score:
- self.run[self._construct_path_with_prefix("model/best_model_score")] = (
- checkpoint_callback.best_model_score.cpu().detach().numpy()
- )
-
- @staticmethod
- def _get_full_model_name(model_path: str, checkpoint_callback: "ReferenceType[ModelCheckpoint]") -> str:
- """Returns model name which is string `model_path` appended to `checkpoint_callback.dirpath`."""
- expected_model_path = f"{checkpoint_callback.dirpath}{os.path.sep}"
- if not model_path.startswith(expected_model_path):
- raise ValueError(f"{model_path} was expected to start with {expected_model_path}.")
- # Remove extension from filepath
- filepath, _ = os.path.splitext(model_path[len(expected_model_path) :])
-
- return filepath
-
- @classmethod
- def _get_full_model_names_from_exp_structure(cls, exp_structure: dict, namespace: str) -> Set[str]:
- """Returns all paths to properties which were already logged in `namespace`"""
- structure_keys = namespace.split(cls.LOGGER_JOIN_CHAR)
- uploaded_models_dict = reduce(lambda d, k: d[k], [exp_structure, *structure_keys])
- return set(cls._dict_paths(uploaded_models_dict))
-
- @classmethod
- def _dict_paths(cls, d: dict, path_in_build: str = None) -> Generator:
- for k, v in d.items():
- path = f"{path_in_build}/{k}" if path_in_build is not None else k
- if not isinstance(v, dict):
- yield path
- else:
- yield from cls._dict_paths(v, path)
-
- @property
- def name(self) -> str:
- """Return the experiment name or 'offline-name' when exp is run in offline mode."""
- return self._run_name
-
- @property
- def version(self) -> str:
- """Return the experiment version.
-
- It's Neptune Run's short_id
- """
- return self._run_short_id
-
- @staticmethod
- def _signal_deprecated_api_usage(f_name, sample_code, raise_exception=False):
- msg_suffix = (
- f"If you are looking for the Neptune logger using legacy Python API,"
- f" it's still available as part of neptune-contrib package:\n"
- f" - https://docs-legacy.neptune.ai/integrations/pytorch_lightning.html\n"
- f"The NeptuneLogger was re-written to use the neptune.new Python API\n"
- f" - https://neptune.ai/blog/neptune-new\n"
- f" - https://docs.neptune.ai/integrations-and-supported-tools/model-training/pytorch-lightning\n"
- f"Instead of `logger.{f_name}` you can use:\n"
- f"\t{sample_code}"
- )
-
- if not raise_exception:
- warnings.warn(
- "The function you've used is deprecated in v1.5.0 and will be removed in v1.7.0. " + msg_suffix
- )
- else:
- raise ValueError("The function you've used is deprecated.\n" + msg_suffix)
-
- @rank_zero_only
- def log_metric(self, metric_name: str, metric_value: Union[torch.Tensor, float, str], step: Optional[int] = None):
- key = f"{self._prefix}/{metric_name}"
- self._signal_deprecated_api_usage("log_metric", f"logger.run['{key}'].log(42)")
- if torch.is_tensor(metric_value):
- metric_value = metric_value.cpu().detach()
-
- self.run[key].log(metric_value, step=step)
-
- @rank_zero_only
- def log_text(self, log_name: str, text: str, step: Optional[int] = None) -> None:
- key = f"{self._prefix}/{log_name}"
- self._signal_deprecated_api_usage("log_text", f"logger.run['{key}].log('text')")
- self.run[key].log(str(text), step=step)
-
- @rank_zero_only
- def log_image(self, log_name: str, image: Union[str, Any], step: Optional[int] = None) -> None:
- key = f"{self._prefix}/{log_name}"
- self._signal_deprecated_api_usage("log_image", f"logger.run['{key}'].log(File('path_to_image'))")
- if isinstance(image, str):
- # if `img` is path to file, convert it to file object
- image = NeptuneFile(image)
- self.run[key].log(image, step=step)
-
- @rank_zero_only
- def log_artifact(self, artifact: str, destination: Optional[str] = None) -> None:
- key = f"{self._prefix}/{self.ARTIFACTS_KEY}/{artifact}"
- self._signal_deprecated_api_usage("log_artifact", f"logger.run['{key}].log('path_to_file')")
- self.run[key].log(destination)
-
- def set_property(self, *args, **kwargs):
- self._signal_deprecated_api_usage(
- "log_artifact", f"logger.run['{self._prefix}/{self.PARAMETERS_KEY}/key'].log(value)", raise_exception=True
- )
-
- def append_tags(self, *args, **kwargs):
- self._signal_deprecated_api_usage(
- "append_tags", "logger.run['sys/tags'].add(['foo', 'bar'])", raise_exception=True
- )
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-TensorBoard Logger
-------------------
-"""
-
-import logging
-import os
-from argparse import Namespace
-from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union
-
-import numpy as np
-import torch
-from torch.utils.tensorboard import SummaryWriter
-from torch.utils.tensorboard.summary import hparams
-
-import pytorch_lightning as pl
-from pytorch_lightning.core.saving import save_hparams_to_yaml
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.cloud_io import get_filesystem
-from pytorch_lightning.utilities.imports import _OMEGACONF_AVAILABLE
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict
-from pytorch_lightning.utilities.logger import _sanitize_params as _utils_sanitize_params
-from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
-
-log = logging.getLogger(__name__)
-
-if _OMEGACONF_AVAILABLE:
- from omegaconf import Container, OmegaConf
-
-
-[docs]class TensorBoardLogger(LightningLoggerBase):
- r"""
- Log to local file system in `TensorBoard <https://www.tensorflow.org/tensorboard>`_ format.
-
- Implemented using :class:`~torch.utils.tensorboard.SummaryWriter`. Logs are saved to
- ``os.path.join(save_dir, name, version)``. This is the default logger in Lightning, it comes
- preinstalled.
-
- Example:
-
- .. testcode::
-
- from pytorch_lightning import Trainer
- from pytorch_lightning.loggers import TensorBoardLogger
-
- logger = TensorBoardLogger("tb_logs", name="my_model")
- trainer = Trainer(logger=logger)
-
- Args:
- save_dir: Save directory
- name: Experiment name. Defaults to ``'default'``. If it is the empty string then no per-experiment
- subdirectory is used.
- version: Experiment version. If version is not specified the logger inspects the save
- directory for existing versions, then automatically assigns the next available version.
- If it is a string then it is used as the run-specific subdirectory name,
- otherwise ``'version_${version}'`` is used.
- log_graph: Adds the computational graph to tensorboard. This requires that
- the user has defined the `self.example_input_array` attribute in their
- model.
- default_hp_metric: Enables a placeholder metric with key `hp_metric` when `log_hyperparams` is
- called without a metric (otherwise calls to log_hyperparams without a metric are ignored).
- prefix: A string to put at the beginning of metric keys.
- sub_dir: Sub-directory to group TensorBoard logs. If a sub_dir argument is passed
- then logs are saved in ``/save_dir/name/version/sub_dir/``. Defaults to ``None`` in which
- logs are saved in ``/save_dir/name/version/``.
- \**kwargs: Additional arguments used by :class:`SummaryWriter` can be passed as keyword
- arguments in this logger. To automatically flush to disk, `max_queue` sets the size
- of the queue for pending logs before flushing. `flush_secs` determines how many seconds
- elapses before flushing.
-
- """
- NAME_HPARAMS_FILE = "hparams.yaml"
- LOGGER_JOIN_CHAR = "-"
-
- def __init__(
- self,
- save_dir: str,
- name: Optional[str] = "lightning_logs",
- version: Optional[Union[int, str]] = None,
- log_graph: bool = False,
- default_hp_metric: bool = True,
- prefix: str = "",
- sub_dir: Optional[str] = None,
- agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,
- agg_default_func: Optional[Callable[[Sequence[float]], float]] = None,
- **kwargs,
- ):
- super().__init__(agg_key_funcs=agg_key_funcs, agg_default_func=agg_default_func)
- self._save_dir = save_dir
- self._name = name or ""
- self._version = version
- self._sub_dir = sub_dir
- self._log_graph = log_graph
- self._default_hp_metric = default_hp_metric
- self._prefix = prefix
- self._fs = get_filesystem(save_dir)
-
- self._experiment = None
- self.hparams = {}
- self._kwargs = kwargs
-
- @property
- def root_dir(self) -> str:
- """Parent directory for all tensorboard checkpoint subdirectories.
-
- If the experiment name parameter is an empty string, no experiment subdirectory is used and the checkpoint will
- be saved in "save_dir/version"
- """
- return os.path.join(self.save_dir, self.name)
-
- @property
- def log_dir(self) -> str:
- """The directory for this run's tensorboard checkpoint.
-
- By default, it is named ``'version_${self.version}'`` but it can be overridden by passing a string value for the
- constructor's version parameter instead of ``None`` or an int.
- """
- # create a pseudo standard path ala test-tube
- version = self.version if isinstance(self.version, str) else f"version_{self.version}"
- log_dir = os.path.join(self.root_dir, version)
- if isinstance(self.sub_dir, str):
- log_dir = os.path.join(log_dir, self.sub_dir)
- log_dir = os.path.expandvars(log_dir)
- log_dir = os.path.expanduser(log_dir)
- return log_dir
-
- @property
- def save_dir(self) -> Optional[str]:
- """Gets the save directory where the TensorBoard experiments are saved.
-
- Returns:
- The local path to the save directory where the TensorBoard experiments are saved.
- """
- return self._save_dir
-
- @property
- def sub_dir(self) -> Optional[str]:
- """Gets the sub directory where the TensorBoard experiments are saved.
-
- Returns:
- The local path to the sub directory where the TensorBoard experiments are saved.
- """
- return self._sub_dir
-
- @property
- @rank_zero_experiment
- def experiment(self) -> SummaryWriter:
- r"""
- Actual tensorboard object. To use TensorBoard features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
-
- Example::
-
- self.logger.experiment.some_tensorboard_function()
-
- """
- if self._experiment is not None:
- return self._experiment
-
- assert rank_zero_only.rank == 0, "tried to init log dirs in non global_rank=0"
- if self.root_dir:
- self._fs.makedirs(self.root_dir, exist_ok=True)
- self._experiment = SummaryWriter(log_dir=self.log_dir, **self._kwargs)
- return self._experiment
-
-[docs] @rank_zero_only
- def log_hyperparams(
- self, params: Union[Dict[str, Any], Namespace], metrics: Optional[Dict[str, Any]] = None
- ) -> None:
- """Record hyperparameters. TensorBoard logs with and without saved hyperparameters are incompatible, the
- hyperparameters are then not displayed in the TensorBoard. Please delete or move the previously saved logs
- to display the new ones with hyperparameters.
-
- Args:
- params: a dictionary-like container with the hyperparameters
- metrics: Dictionary with metric names as keys and measured quantities as values
- """
-
- params = _convert_params(params)
-
- # store params to output
- if _OMEGACONF_AVAILABLE and isinstance(params, Container):
- self.hparams = OmegaConf.merge(self.hparams, params)
- else:
- self.hparams.update(params)
-
- # format params into the suitable for tensorboard
- params = _flatten_dict(params)
- params = self._sanitize_params(params)
-
- if metrics is None:
- if self._default_hp_metric:
- metrics = {"hp_metric": -1}
- elif not isinstance(metrics, dict):
- metrics = {"hp_metric": metrics}
-
- if metrics:
- self.log_metrics(metrics, 0)
- exp, ssi, sei = hparams(params, metrics)
- writer = self.experiment._get_file_writer()
- writer.add_summary(exp)
- writer.add_summary(ssi)
- writer.add_summary(sei)
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
- assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
-
- metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
-
- for k, v in metrics.items():
- if isinstance(v, torch.Tensor):
- v = v.item()
-
- if isinstance(v, dict):
- self.experiment.add_scalars(k, v, step)
- else:
- try:
- self.experiment.add_scalar(k, v, step)
- # todo: specify the possible exception
- except Exception as ex:
- m = f"\n you tried to log {v} which is currently not supported. Try a dict or a scalar/tensor."
- raise ValueError(m) from ex
-
-[docs] @rank_zero_only
- def log_graph(self, model: "pl.LightningModule", input_array=None):
- if self._log_graph:
- if input_array is None:
- input_array = model.example_input_array
-
- if input_array is not None:
- input_array = model._apply_batch_transfer_handler(input_array)
- model._running_torchscript = True
- self.experiment.add_graph(model, input_array)
- model._running_torchscript = False
- else:
- rank_zero_warn(
- "Could not log computational graph since the"
- " `model.example_input_array` attribute is not set"
- " or `input_array` was not given",
- )
-
-[docs] @rank_zero_only
- def save(self) -> None:
- super().save()
- dir_path = self.log_dir
-
- # prepare the file path
- hparams_file = os.path.join(dir_path, self.NAME_HPARAMS_FILE)
-
- # save the metatags file if it doesn't exist and the log directory exists
- if self._fs.isdir(dir_path) and not self._fs.isfile(hparams_file):
- save_hparams_to_yaml(hparams_file, self.hparams)
-
-[docs] @rank_zero_only
- def finalize(self, status: str) -> None:
- self.experiment.flush()
- self.experiment.close()
- self.save()
-
- @property
- def name(self) -> str:
- """Get the name of the experiment.
-
- Returns:
- The name of the experiment.
- """
- return self._name
-
- @property
- def version(self) -> int:
- """Get the experiment version.
-
- Returns:
- The experiment version if specified else the next version.
- """
- if self._version is None:
- self._version = self._get_next_version()
- return self._version
-
- def _get_next_version(self):
- root_dir = self.root_dir
-
- try:
- listdir_info = self._fs.listdir(root_dir)
- except OSError:
- log.warning("Missing logger folder: %s", root_dir)
- return 0
-
- existing_versions = []
- for listing in listdir_info:
- d = listing["name"]
- bn = os.path.basename(d)
- if self._fs.isdir(d) and bn.startswith("version_"):
- dir_ver = bn.split("_")[1].replace("/", "")
- existing_versions.append(int(dir_ver))
- if len(existing_versions) == 0:
- return 0
-
- return max(existing_versions) + 1
-
- @staticmethod
- def _sanitize_params(params: Dict[str, Any]) -> Dict[str, Any]:
- params = _utils_sanitize_params(params)
- # logging of arrays with dimension > 1 is not supported, sanitize as string
- return {k: str(v) if isinstance(v, (torch.Tensor, np.ndarray)) and v.ndim > 1 else v for k, v in params.items()}
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["_experiment"] = None
- return state
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Weights and Biases Logger
--------------------------
-"""
-import os
-from argparse import Namespace
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union
-from weakref import ReferenceType
-
-import torch.nn as nn
-
-from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
-from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.imports import _WANDB_GREATER_EQUAL_0_10_22, _WANDB_GREATER_EQUAL_0_12_10
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict, _sanitize_callable_params
-from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
-
-try:
- import wandb
- from wandb.wandb_run import Run
-except ModuleNotFoundError:
- # needed for test mocks, these tests shall be updated
- wandb, Run = None, None
-
-
-[docs]class WandbLogger(LightningLoggerBase):
- r"""
- Log using `Weights and Biases <https://docs.wandb.ai/integrations/lightning>`_.
-
- **Installation and set-up**
-
- Install with pip:
-
- .. code-block:: bash
-
- pip install wandb
-
- Create a `WandbLogger` instance:
-
- .. code-block:: python
-
- from pytorch_lightning.loggers import WandbLogger
-
- wandb_logger = WandbLogger(project="MNIST")
-
- Pass the logger instance to the `Trainer`:
-
- .. code-block:: python
-
- trainer = Trainer(logger=wandb_logger)
-
- A new W&B run will be created when training starts if you have not created one manually before with `wandb.init()`.
-
- **Log metrics**
-
- Log from :class:`~pytorch_lightning.core.lightning.LightningModule`:
-
- .. code-block:: python
-
- class LitModule(LightningModule):
- def training_step(self, batch, batch_idx):
- self.log("train/loss", loss)
-
- Use directly wandb module:
-
- .. code-block:: python
-
- wandb.log({"train/loss": loss})
-
- **Log hyper-parameters**
-
- Save :class:`~pytorch_lightning.core.lightning.LightningModule` parameters:
-
- .. code-block:: python
-
- class LitModule(LightningModule):
- def __init__(self, *args, **kwarg):
- self.save_hyperparameters()
-
- Add other config parameters:
-
- .. code-block:: python
-
- # add one parameter
- wandb_logger.experiment.config["key"] = value
-
- # add multiple parameters
- wandb_logger.experiment.config.update({key1: val1, key2: val2})
-
- # use directly wandb module
- wandb.config["key"] = value
- wandb.config.update()
-
- **Log gradients, parameters and model topology**
-
- Call the `watch` method for automatically tracking gradients:
-
- .. code-block:: python
-
- # log gradients and model topology
- wandb_logger.watch(model)
-
- # log gradients, parameter histogram and model topology
- wandb_logger.watch(model, log="all")
-
- # change log frequency of gradients and parameters (100 steps by default)
- wandb_logger.watch(model, log_freq=500)
-
- # do not log graph (in case of errors)
- wandb_logger.watch(model, log_graph=False)
-
- The `watch` method adds hooks to the model which can be removed at the end of training:
-
- .. code-block:: python
-
- wandb_logger.unwatch(model)
-
- **Log model checkpoints**
-
- Log model checkpoints at the end of training:
-
- .. code-block:: python
-
- wandb_logger = WandbLogger(log_model=True)
-
- Log model checkpoints as they get created during training:
-
- .. code-block:: python
-
- wandb_logger = WandbLogger(log_model="all")
-
- Custom checkpointing can be set up through :class:`~pytorch_lightning.callbacks.ModelCheckpoint`:
-
- .. code-block:: python
-
- # log model only if `val_accuracy` increases
- wandb_logger = WandbLogger(log_model="all")
- checkpoint_callback = ModelCheckpoint(monitor="val_accuracy", mode="max")
- trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback])
-
- `latest` and `best` aliases are automatically set to easily retrieve a model checkpoint:
-
- .. code-block:: python
-
- # reference can be retrieved in artifacts panel
- # "VERSION" can be a version (ex: "v2") or an alias ("latest or "best")
- checkpoint_reference = "USER/PROJECT/MODEL-RUN_ID:VERSION"
-
- # download checkpoint locally (if not already cached)
- run = wandb.init(project="MNIST")
- artifact = run.use_artifact(checkpoint_reference, type="model")
- artifact_dir = artifact.download()
-
- # load checkpoint
- model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt")
-
- **Log media**
-
- Log text with:
-
- .. code-block:: python
-
- # using columns and data
- columns = ["input", "label", "prediction"]
- data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]]
- wandb_logger.log_text(key="samples", columns=columns, data=data)
-
- # using a pandas DataFrame
- wandb_logger.log_text(key="samples", dataframe=my_dataframe)
-
- Log images with:
-
- .. code-block:: python
-
- # using tensors, numpy arrays or PIL images
- wandb_logger.log_image(key="samples", images=[img1, img2])
-
- # adding captions
- wandb_logger.log_image(key="samples", images=[img1, img2], caption=["tree", "person"])
-
- # using file path
- wandb_logger.log_image(key="samples", images=["img_1.jpg", "img_2.jpg"])
-
- More arguments can be passed for logging segmentation masks and bounding boxes. Refer to
- `Image Overlays documentation <https://docs.wandb.ai/guides/track/log/media#image-overlays>`_.
-
- **Log Tables**
-
- `W&B Tables <https://docs.wandb.ai/guides/data-vis>`_ can be used to log, query and analyze tabular data.
-
- They support any type of media (text, image, video, audio, molecule, html, etc) and are great for storing,
- understanding and sharing any form of data, from datasets to model predictions.
-
- .. code-block:: python
-
- columns = ["caption", "image", "sound"]
- data = [["cheese", wandb.Image(img_1), wandb.Audio(snd_1)], ["wine", wandb.Image(img_2), wandb.Audio(snd_2)]]
- wandb_logger.log_table(key="samples", columns=columns, data=data)
-
- See Also:
- - `Demo in Google Colab <http://wandb.me/lightning>`__ with hyperparameter search and model logging
- - `W&B Documentation <https://docs.wandb.ai/integrations/lightning>`__
-
- Args:
- name: Display name for the run.
- save_dir: Path where data is saved (wandb dir by default).
- offline: Run offline (data can be streamed later to wandb servers).
- id: Sets the version, mainly used to resume a previous run.
- version: Same as id.
- anonymous: Enables or explicitly disables anonymous logging.
- project: The name of the project to which this run will belong.
- log_model: Log checkpoints created by :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint`
- as W&B artifacts. `latest` and `best` aliases are automatically set.
-
- * if ``log_model == 'all'``, checkpoints are logged during training.
- * if ``log_model == True``, checkpoints are logged at the end of training, except when
- :paramref:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint.save_top_k` ``== -1``
- which also logs every checkpoint during training.
- * if ``log_model == False`` (default), no checkpoint is logged.
-
- prefix: A string to put at the beginning of metric keys.
- experiment: WandB experiment object. Automatically set when creating a run.
- \**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc.
-
- Raises:
- ModuleNotFoundError:
- If required WandB package is not installed on the device.
- MisconfigurationException:
- If both ``log_model`` and ``offline`` is set to ``True``.
-
- """
-
- LOGGER_JOIN_CHAR = "-"
-
- def __init__(
- self,
- name: Optional[str] = None,
- save_dir: Optional[str] = None,
- offline: Optional[bool] = False,
- id: Optional[str] = None,
- anonymous: Optional[bool] = None,
- version: Optional[str] = None,
- project: Optional[str] = None,
- log_model: Union[str, bool] = False,
- experiment=None,
- prefix: Optional[str] = "",
- agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,
- agg_default_func: Optional[Callable[[Sequence[float]], float]] = None,
- **kwargs,
- ):
- if wandb is None:
- raise ModuleNotFoundError(
- "You want to use `wandb` logger which is not installed yet,"
- " install it with `pip install wandb`." # pragma: no-cover
- )
-
- if offline and log_model:
- raise MisconfigurationException(
- f"Providing log_model={log_model} and offline={offline} is an invalid configuration"
- " since model checkpoints cannot be uploaded in offline mode.\n"
- "Hint: Set `offline=False` to log your model."
- )
-
- if log_model and not _WANDB_GREATER_EQUAL_0_10_22:
- rank_zero_warn(
- f"Providing log_model={log_model} requires wandb version >= 0.10.22"
- " for logging associated model metadata.\n"
- "Hint: Upgrade with `pip install --upgrade wandb`."
- )
-
- super().__init__(agg_key_funcs=agg_key_funcs, agg_default_func=agg_default_func)
- self._offline = offline
- self._log_model = log_model
- self._prefix = prefix
- self._experiment = experiment
- self._logged_model_time = {}
- self._checkpoint_callback = None
- # set wandb init arguments
- anonymous_lut = {True: "allow", False: None}
- self._wandb_init = dict(
- name=name,
- project=project,
- id=version or id,
- dir=save_dir,
- resume="allow",
- anonymous=anonymous_lut.get(anonymous, anonymous),
- )
- self._wandb_init.update(**kwargs)
- # extract parameters
- self._save_dir = self._wandb_init.get("dir")
- self._name = self._wandb_init.get("name")
- self._id = self._wandb_init.get("id")
- # start wandb run (to create an attach_id for distributed modes)
- if _WANDB_GREATER_EQUAL_0_12_10:
- wandb.require("service")
- _ = self.experiment
-
- def __getstate__(self):
- state = self.__dict__.copy()
- # args needed to reload correct experiment
- if self._experiment is not None:
- state["_id"] = getattr(self._experiment, "id", None)
- state["_attach_id"] = getattr(self._experiment, "_attach_id", None)
- state["_name"] = self._experiment.project_name()
-
- # cannot be pickled
- state["_experiment"] = None
- return state
-
- @property
- @rank_zero_experiment
- def experiment(self) -> Run:
- r"""
-
- Actual wandb object. To use wandb features in your
- :class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
-
- Example::
-
- .. code-block:: python
-
- self.logger.experiment.some_wandb_function()
-
- """
- if self._experiment is None:
- if self._offline:
- os.environ["WANDB_MODE"] = "dryrun"
-
- attach_id = getattr(self, "_attach_id", None)
- if wandb.run is not None:
- # wandb process already created in this instance
- rank_zero_warn(
- "There is a wandb run already in progress and newly created instances of `WandbLogger` will reuse"
- " this run. If this is not desired, call `wandb.finish()` before instantiating `WandbLogger`."
- )
- self._experiment = wandb.run
- elif attach_id is not None and hasattr(wandb, "_attach"):
- # attach to wandb process referenced
- self._experiment = wandb._attach(attach_id)
- else:
- # create new wandb process
- self._experiment = wandb.init(**self._wandb_init)
-
- # define default x-axis
- if getattr(self._experiment, "define_metric", None):
- self._experiment.define_metric("trainer/global_step")
- self._experiment.define_metric("*", step_metric="trainer/global_step", step_sync=True)
-
- return self._experiment
-
- def watch(self, model: nn.Module, log: str = "gradients", log_freq: int = 100, log_graph: bool = True):
- self.experiment.watch(model, log=log, log_freq=log_freq, log_graph=log_graph)
-
-[docs] @rank_zero_only
- def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
- params = _convert_params(params)
- params = _flatten_dict(params)
- params = _sanitize_callable_params(params)
- self.experiment.config.update(params, allow_val_change=True)
-
-[docs] @rank_zero_only
- def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
- assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
-
- metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
- if step is not None:
- self.experiment.log({**metrics, "trainer/global_step": step})
- else:
- self.experiment.log(metrics)
-
-[docs] @rank_zero_only
- def log_table(
- self,
- key: str,
- columns: List[str] = None,
- data: List[List[Any]] = None,
- dataframe: Any = None,
- step: Optional[int] = None,
- ) -> None:
- """Log a Table containing any object type (text, image, audio, video, molecule, html, etc).
-
- Can be defined either with `columns` and `data` or with `dataframe`.
- """
-
- metrics = {key: wandb.Table(columns=columns, data=data, dataframe=dataframe)}
- self.log_metrics(metrics, step)
-
-[docs] @rank_zero_only
- def log_text(
- self,
- key: str,
- columns: List[str] = None,
- data: List[List[str]] = None,
- dataframe: Any = None,
- step: Optional[int] = None,
- ) -> None:
- """Log text as a Table.
-
- Can be defined either with `columns` and `data` or with `dataframe`.
- """
-
- self.log_table(key, columns, data, dataframe, step)
-
-[docs] @rank_zero_only
- def log_image(self, key: str, images: List[Any], step: Optional[int] = None, **kwargs: str) -> None:
- """Log images (tensors, numpy arrays, PIL Images or file paths).
-
- Optional kwargs are lists passed to each image (ex: caption, masks, boxes).
- """
- if not isinstance(images, list):
- raise TypeError(f'Expected a list as "images", found {type(images)}')
- n = len(images)
- for k, v in kwargs.items():
- if len(v) != n:
- raise ValueError(f"Expected {n} items but only found {len(v)} for {k}")
- kwarg_list = [{k: kwargs[k][i] for k in kwargs.keys()} for i in range(n)]
- metrics = {key: [wandb.Image(img, **kwarg) for img, kwarg in zip(images, kwarg_list)]}
- self.log_metrics(metrics, step)
-
- @property
- def save_dir(self) -> Optional[str]:
- """Gets the save directory.
-
- Returns:
- The path to the save directory.
- """
- return self._save_dir
-
- @property
- def name(self) -> Optional[str]:
- """Gets the name of the experiment.
-
- Returns:
- The name of the experiment if the experiment exists else the name given to the constructor.
- """
- # don't create an experiment if we don't have one
- return self._experiment.project_name() if self._experiment else self._name
-
- @property
- def version(self) -> Optional[str]:
- """Gets the id of the experiment.
-
- Returns:
- The id of the experiment if the experiment exists else the id given to the constructor.
- """
- # don't create an experiment if we don't have one
- return self._experiment.id if self._experiment else self._id
-
-[docs] def after_save_checkpoint(self, checkpoint_callback: "ReferenceType[ModelCheckpoint]") -> None:
- # log checkpoints as artifacts
- if self._log_model == "all" or self._log_model is True and checkpoint_callback.save_top_k == -1:
- self._scan_and_log_checkpoints(checkpoint_callback)
- elif self._log_model is True:
- self._checkpoint_callback = checkpoint_callback
-
-[docs] @rank_zero_only
- def finalize(self, status: str) -> None:
- # log checkpoints as artifacts
- if self._checkpoint_callback:
- self._scan_and_log_checkpoints(self._checkpoint_callback)
-
- def _scan_and_log_checkpoints(self, checkpoint_callback: "ReferenceType[ModelCheckpoint]") -> None:
- # get checkpoints to be saved with associated score
- checkpoints = {
- checkpoint_callback.last_model_path: checkpoint_callback.current_score,
- checkpoint_callback.best_model_path: checkpoint_callback.best_model_score,
- **checkpoint_callback.best_k_models,
- }
- checkpoints = sorted((Path(p).stat().st_mtime, p, s) for p, s in checkpoints.items() if Path(p).is_file())
- checkpoints = [
- c for c in checkpoints if c[1] not in self._logged_model_time.keys() or self._logged_model_time[c[1]] < c[0]
- ]
-
- # log iteratively all new checkpoints
- for t, p, s in checkpoints:
- metadata = (
- {
- "score": s,
- "original_filename": Path(p).name,
- "ModelCheckpoint": {
- k: getattr(checkpoint_callback, k)
- for k in [
- "monitor",
- "mode",
- "save_last",
- "save_top_k",
- "save_weights_only",
- "_every_n_train_steps",
- ]
- # ensure it does not break if `ModelCheckpoint` args change
- if hasattr(checkpoint_callback, k)
- },
- }
- if _WANDB_GREATER_EQUAL_0_10_22
- else None
- )
- artifact = wandb.Artifact(name=f"model-{self.experiment.id}", type="model", metadata=metadata)
- artifact.add_file(p, name="model.ckpt")
- aliases = ["latest", "best"] if p == checkpoint_callback.best_model_path else ["latest"]
- self.experiment.log_artifact(artifact, aliases=aliases)
- # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name)
- self._logged_model_time[p] = t
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import inspect
-from abc import ABC, abstractmethod
-from typing import Any, Dict, Generic, Optional, Type, TypeVar, Union
-
-from deprecate import void
-from torchmetrics import Metric
-
-import pytorch_lightning as pl
-from pytorch_lightning.trainer.connectors.logger_connector.result import _ResultCollection
-from pytorch_lightning.trainer.progress import BaseProgress
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.imports import _fault_tolerant_training
-
-T = TypeVar("T") # the output type of `run`
-
-
-class Loop(ABC, Generic[T]):
- """Basic Loops interface. All classes derived from this must implement the following properties and methods:
-
- * :attr:`done` (property): Condition to break the loop
- * :attr:`reset` (method): Resets the internal state between multiple calls of :attr:`run`
- * :attr:`advance` (method): Implements one step of the loop
-
- This class implements the following loop structure:
-
- .. code-block:: python
-
- on_run_start()
-
- while not done:
- on_advance_start()
- advance()
- on_advance_end()
-
- on_run_end()
- """
-
- def __init__(self) -> None:
- self._restarting = False
- self._trainer: Optional["pl.Trainer"] = None
-
- @property
- def trainer(self) -> "pl.Trainer":
- if self._trainer is None:
- raise RuntimeError("The loop is not attached to a Trainer.")
- return self._trainer
-
- @trainer.setter
- def trainer(self, trainer: "pl.Trainer") -> None:
- """Connects this loop's trainer and its children."""
- if not isinstance(trainer, pl.Trainer):
- raise MisconfigurationException(
- f"Loop {self.__class__.__name__} should be connected to a `Trainer`, found: {trainer}."
- )
- self._trainer = trainer
- for v in self.__dict__.values():
- if isinstance(v, Loop):
- v.trainer = trainer
-
- @property
- def restarting(self) -> bool:
- """Whether the state of this loop was reloaded and it needs to restart."""
- return self._restarting
-
- @restarting.setter
- def restarting(self, restarting: bool) -> None:
- """Connects this loop's restarting value and its children."""
- self._restarting = restarting
- for loop in vars(self).values():
- if isinstance(loop, Loop):
- loop.restarting = restarting
-
- @property
- @abstractmethod
- def done(self) -> bool:
- """Property indicating when the loop is finished.
-
- Example::
-
- @property
- def done(self):
- return self.trainer.global_step >= self.trainer.max_steps
- """
-
- @property
- def skip(self) -> bool:
- """Determine whether to return immediately from the call to :meth:`run`.
-
- Example::
-
- @property
- def skip(self):
- return len(self.trainer.train_dataloader) == 0
- """
- return False
-
- def connect(self, **kwargs: "Loop") -> None:
- """Optionally connect one or multiple loops to this one.
-
- Linked loops should form a tree.
- """
-
- def replace(self, **loops: Union["Loop", Type["Loop"]]) -> None:
- """Optionally replace one or multiple of this loop's sub-loops.
-
- This methods takes care of instantiating the class (if necessary) with all existing arguments, connecting all
- sub-loops of the old loop to the new instance, setting the ``Trainer`` reference, and connecting the new loop to
- the parent.
-
- Args:
- **loops: ``Loop`` subclasses or instances. The name used should match the loop attribute name you want to
- replace.
-
- Raises:
- MisconfigurationException: When passing a ``Loop`` class, if the ``__init__`` arguments do not match those
- of the Loop class it replaces.
- """
- new_loops = {}
-
- for name, type_or_object in loops.items():
- old_loop = getattr(self, name)
-
- if isinstance(type_or_object, type):
- # compare the signatures
- old_parameters = inspect.signature(old_loop.__class__.__init__).parameters
- current_parameters = inspect.signature(type_or_object.__init__).parameters
- if old_parameters != current_parameters:
- raise MisconfigurationException(
- f"`{self.__class__.__name__}.replace({type_or_object.__name__})` can only be used if the"
- f" `__init__` signatures match but `{old_loop.__class__.__name__}` does not."
- )
- # instantiate the loop
- kwargs = {p: getattr(old_loop, p) for p in old_parameters if p != "self"}
- loop = type_or_object(**kwargs)
- else:
- loop = type_or_object
-
- # connect sub-loops
- kwargs = {n: l for n, l in old_loop.__dict__.items() if isinstance(l, Loop)}
- loop.connect(**kwargs)
- # set the trainer reference
- loop.trainer = self.trainer
-
- new_loops[name] = loop
- # connect to self
- self.connect(**new_loops)
-
- def on_skip(self) -> T:
- """The function to run when :meth:`run` should be skipped, determined by the condition in :attr:`skip`.
-
- Returns:
- the default output value of :meth:`on_run_end`
- """
-
-[docs] def run(self, *args: Any, **kwargs: Any) -> T:
- """The main entry point to the loop.
-
- Will frequently check the :attr:`done` condition and calls :attr:`advance`
- until :attr:`done` evaluates to ``True``.
-
- Override this if you wish to change the default behavior. The default implementation is:
-
- Example::
-
- def run(self, *args, **kwargs):
- if self.skip:
- return self.on_skip()
-
- self.reset()
- self.on_run_start(*args, **kwargs)
-
- while not self.done:
- self.advance(*args, **kwargs)
-
- output = self.on_run_end()
- return output
-
- Returns:
- The output of :attr:`on_run_end` (often outputs collected from each step of the loop)
- """
- if self.skip:
- return self.on_skip()
-
- self.reset()
-
- self.on_run_start(*args, **kwargs)
-
- while not self.done:
- try:
- self.on_advance_start(*args, **kwargs)
- self.advance(*args, **kwargs)
- self.on_advance_end()
- self._restarting = False
- except StopIteration:
- break
- self._restarting = False
-
- output = self.on_run_end()
- return output
-
-[docs] @abstractmethod
- def reset(self) -> None:
- """Resets the internal state of the loop at the beginning of each call to :attr:`run`.
-
- Example::
-
- def reset(self):
- # reset your internal state or add custom logic
- # if you expect run() to be called multiple times
- self.current_iteration = 0
- self.outputs = []
- """
-
- def on_run_start(self, *args: Any, **kwargs: Any) -> None:
- """Hook to be called as the first thing after entering :attr:`run` (except the state reset).
-
- Accepts all arguments passed to :attr:`run`.
- """
- void(*args, **kwargs)
-
- def on_advance_start(self, *args: Any, **kwargs: Any) -> None:
- """Hook to be called each time before :attr:`advance` is called.
-
- Accepts all arguments passed to :attr`run`.
- """
- void(*args, **kwargs)
-
-[docs] @abstractmethod
- def advance(self, *args: Any, **kwargs: Any) -> None:
- """Performs a single step.
-
- Accepts all arguments passed to :attr:`run`.
-
- Example::
-
- def advance(self, iterator):
- batch = next(iterator)
- loss = self.trainer.lightning_module.training_step(batch, batch_idx)
- ...
- """
-
- def on_advance_end(self) -> None:
- """Hook to be called each time after :attr:`advance` is called."""
-
- def on_run_end(self) -> T:
- """Hook to be called at the end of the run.
-
- Its return argument is returned from :attr:`run`.
- """
-
- def teardown(self) -> None:
- """Use to release memory etc."""
-
- def on_save_checkpoint(self) -> Dict:
- """Called when saving a model checkpoint, use to persist loop state.
-
- Returns:
- The current loop state.
- """
- return {}
-
- def on_load_checkpoint(self, state_dict: Dict) -> None:
- """Called when loading a model checkpoint, use to reload loop state."""
-
- def state_dict(self, destination: Optional[Dict] = None, prefix: str = "") -> Dict:
- """The state dict is determined by the state and progress of this loop and all its children.
-
- Args:
- destination: An existing dictionary to update with this loop's state. By default a new dictionary
- is returned.
- prefix: A prefix for each key in the state dictionary
- """
- if destination is None:
- destination = {}
-
- destination[prefix + "state_dict"] = self.on_save_checkpoint()
-
- # do not get the mode from `self.trainer` because it might not have been attached yet
- ft_enabled = _fault_tolerant_training()
- for k, v in self.__dict__.items():
- key = prefix + k
- if isinstance(v, BaseProgress):
- destination[key] = v.state_dict()
- elif isinstance(v, Loop):
- v.state_dict(destination, key + ".")
- elif ft_enabled and isinstance(v, _ResultCollection):
- # sync / unsync metrics
- v.sync()
- destination[key] = v.state_dict()
- v.unsync()
-
- return destination
-
- def load_state_dict(
- self,
- state_dict: Dict,
- prefix: str = "",
- metrics: Optional[Dict[str, Metric]] = None,
- ) -> None:
- """Loads the state of this loop and all its children."""
- self._load_from_state_dict(state_dict.copy(), prefix, metrics)
- for k, v in self.__dict__.items():
- if isinstance(v, Loop):
- v.load_state_dict(state_dict.copy(), prefix + k + ".")
- self.restarting = True
-
- def _load_from_state_dict(self, state_dict: Dict, prefix: str, metrics: Optional[Dict[str, Metric]] = None) -> None:
- for k, v in self.__dict__.items():
- key = prefix + k
- if key not in state_dict:
- # compatibility with old checkpoints
- continue
-
- if isinstance(v, BaseProgress):
- v.load_state_dict(state_dict[key])
- elif (
- isinstance(v, _ResultCollection)
- and self.trainer is not None
- and self.trainer.lightning_module is not None
- ):
- metric_attributes = {
- name: module
- for name, module in self.trainer.lightning_module.named_modules()
- if isinstance(module, Metric)
- }
- if metrics:
- metric_attributes.update(metrics)
-
- # The `_ResultCollection` objects have 2 types of metrics: `Tensor` and `torchmetrics.Metric`.
- # When creating a checkpoint, the `Metric`s are dropped from the loop `state_dict` to serialize only
- # Python primitives. However, their states are saved with the model's `state_dict`.
- # On reload, we need to re-attach the `Metric`s back to the `_ResultCollection`.
- # The references are provided through the `metric_attributes` dictionary.
- v.load_state_dict(state_dict[key], metrics=metric_attributes, sync_fn=self.trainer.strategy.reduce)
-
- if not self.trainer.is_global_zero:
- v.reset(metrics=False)
-
- if prefix + "state_dict" in state_dict: # compatibility with old checkpoints
- self.on_load_checkpoint(state_dict[prefix + "state_dict"])
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Trainer to automate the training."""
-import inspect
-import logging
-import math
-import os
-import traceback
-import warnings
-from argparse import ArgumentParser, Namespace
-from copy import deepcopy
-from datetime import timedelta
-from pathlib import Path
-from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Type, Union
-from weakref import proxy
-
-import torch
-from packaging.version import Version
-from torch.optim import Optimizer
-from torch.utils.data import DataLoader
-
-import pytorch_lightning as pl
-from pytorch_lightning.accelerators import Accelerator, GPUAccelerator, HPUAccelerator, IPUAccelerator, TPUAccelerator
-from pytorch_lightning.callbacks import Callback, EarlyStopping, ModelCheckpoint, ProgressBarBase
-from pytorch_lightning.callbacks.prediction_writer import BasePredictionWriter
-from pytorch_lightning.core.datamodule import LightningDataModule
-from pytorch_lightning.core.optimizer import LightningOptimizer
-from pytorch_lightning.loggers import LightningLoggerBase
-from pytorch_lightning.loggers.base import DummyLogger, LoggerCollection
-from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
-from pytorch_lightning.loops import PredictionLoop, TrainingEpochLoop
-from pytorch_lightning.loops.dataloader.evaluation_loop import EvaluationLoop
-from pytorch_lightning.loops.fit_loop import FitLoop
-from pytorch_lightning.loops.utilities import _parse_loop_limits, _reset_progress
-from pytorch_lightning.plugins import (
- ApexMixedPrecisionPlugin,
- NativeMixedPrecisionPlugin,
- PLUGIN_INPUT,
- PrecisionPlugin,
-)
-from pytorch_lightning.plugins.environments.slurm_environment import SLURMEnvironment
-from pytorch_lightning.profiler import (
- AdvancedProfiler,
- BaseProfiler,
- PassThroughProfiler,
- Profiler,
- PyTorchProfiler,
- SimpleProfiler,
- XLAProfiler,
-)
-from pytorch_lightning.strategies import ParallelStrategy, Strategy
-from pytorch_lightning.strategies.ddp_spawn import DDPSpawnStrategy
-from pytorch_lightning.trainer.callback_hook import TrainerCallbackHookMixin
-from pytorch_lightning.trainer.configuration_validator import verify_loop_configurations
-from pytorch_lightning.trainer.connectors.accelerator_connector import AcceleratorConnector
-from pytorch_lightning.trainer.connectors.callback_connector import CallbackConnector
-from pytorch_lightning.trainer.connectors.checkpoint_connector import CheckpointConnector
-from pytorch_lightning.trainer.connectors.data_connector import DataConnector
-from pytorch_lightning.trainer.connectors.logger_connector import LoggerConnector
-from pytorch_lightning.trainer.connectors.logger_connector.result import _ResultCollection
-from pytorch_lightning.trainer.connectors.signal_connector import SignalConnector
-from pytorch_lightning.trainer.data_loading import TrainerDataLoadingMixin
-from pytorch_lightning.trainer.optimizers import TrainerOptimizersMixin
-from pytorch_lightning.trainer.states import RunningStage, TrainerFn, TrainerState, TrainerStatus
-from pytorch_lightning.trainer.supporters import CombinedLoader
-from pytorch_lightning.tuner.lr_finder import _LRFinder
-from pytorch_lightning.tuner.tuning import Tuner
-from pytorch_lightning.utilities import (
- _HPU_AVAILABLE,
- _IPU_AVAILABLE,
- _TPU_AVAILABLE,
- AMPType,
- GradClipAlgorithmType,
- parsing,
-)
-from pytorch_lightning.utilities.apply_func import apply_to_collection
-from pytorch_lightning.utilities.argparse import (
- _defaults_from_env_vars,
- add_argparse_args,
- from_argparse_args,
- parse_argparser,
- parse_env_variables,
-)
-from pytorch_lightning.utilities.auto_restart import _add_capture_metadata_collate
-from pytorch_lightning.utilities.cloud_io import get_filesystem
-from pytorch_lightning.utilities.data import _auto_add_worker_init_fn, has_len_all_ranks
-from pytorch_lightning.utilities.distributed import distributed_available
-from pytorch_lightning.utilities.exceptions import ExitGracefullyException, MisconfigurationException
-from pytorch_lightning.utilities.imports import _fault_tolerant_training
-from pytorch_lightning.utilities.meta import is_on_meta_device, materialize_module
-from pytorch_lightning.utilities.model_helpers import is_overridden
-from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_info, rank_zero_warn
-from pytorch_lightning.utilities.seed import isolate_rng
-from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
-from pytorch_lightning.utilities.types import (
- _EVALUATE_OUTPUT,
- _PATH,
- _PREDICT_OUTPUT,
- EVAL_DATALOADERS,
- LRSchedulerConfig,
- STEP_OUTPUT,
- TRAIN_DATALOADERS,
-)
-from pytorch_lightning.utilities.warnings import PossibleUserWarning
-
-log = logging.getLogger(__name__)
-# warnings to ignore in trainer
-warnings.filterwarnings(
- "ignore", message="torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead"
-)
-
-
-class Trainer(
- TrainerCallbackHookMixin, # TODO: Remove in v1.8
- TrainerOptimizersMixin, # TODO: Remove in v1.8
- TrainerDataLoadingMixin, # TODO: Remove in v1.8
-):
-[docs] @_defaults_from_env_vars
- def __init__(
- self,
- logger: Union[LightningLoggerBase, Iterable[LightningLoggerBase], bool] = True,
- checkpoint_callback: Optional[bool] = None,
- enable_checkpointing: bool = True,
- callbacks: Optional[Union[List[Callback], Callback]] = None,
- default_root_dir: Optional[str] = None,
- gradient_clip_val: Optional[Union[int, float]] = None,
- gradient_clip_algorithm: Optional[str] = None,
- process_position: int = 0,
- num_nodes: int = 1,
- num_processes: Optional[int] = None,
- devices: Optional[Union[List[int], str, int]] = None,
- gpus: Optional[Union[List[int], str, int]] = None,
- auto_select_gpus: bool = False,
- tpu_cores: Optional[Union[List[int], str, int]] = None,
- ipus: Optional[int] = None,
- log_gpu_memory: Optional[str] = None, # TODO: Remove in 1.7
- progress_bar_refresh_rate: Optional[int] = None, # TODO: remove in v1.7
- enable_progress_bar: bool = True,
- overfit_batches: Union[int, float] = 0.0,
- track_grad_norm: Union[int, float, str] = -1,
- check_val_every_n_epoch: int = 1,
- fast_dev_run: Union[int, bool] = False,
- accumulate_grad_batches: Optional[Union[int, Dict[int, int]]] = None,
- max_epochs: Optional[int] = None,
- min_epochs: Optional[int] = None,
- max_steps: int = -1,
- min_steps: Optional[int] = None,
- max_time: Optional[Union[str, timedelta, Dict[str, int]]] = None,
- limit_train_batches: Optional[Union[int, float]] = None,
- limit_val_batches: Optional[Union[int, float]] = None,
- limit_test_batches: Optional[Union[int, float]] = None,
- limit_predict_batches: Optional[Union[int, float]] = None,
- val_check_interval: Optional[Union[int, float]] = None,
- flush_logs_every_n_steps: Optional[int] = None,
- log_every_n_steps: int = 50,
- accelerator: Optional[Union[str, Accelerator]] = None,
- strategy: Optional[Union[str, Strategy]] = None,
- sync_batchnorm: bool = False,
- precision: Union[int, str] = 32,
- enable_model_summary: bool = True,
- weights_summary: Optional[str] = "top",
- weights_save_path: Optional[str] = None, # TODO: Remove in 1.8
- num_sanity_val_steps: int = 2,
- resume_from_checkpoint: Optional[Union[Path, str]] = None,
- profiler: Optional[Union[BaseProfiler, str]] = None,
- benchmark: Optional[bool] = None,
- deterministic: bool = False,
- reload_dataloaders_every_n_epochs: int = 0,
- auto_lr_find: Union[bool, str] = False,
- replace_sampler_ddp: bool = True,
- detect_anomaly: bool = False,
- auto_scale_batch_size: Union[str, bool] = False,
- prepare_data_per_node: Optional[bool] = None,
- plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]] = None,
- amp_backend: str = "native",
- amp_level: Optional[str] = None,
- move_metrics_to_cpu: bool = False,
- multiple_trainloader_mode: str = "max_size_cycle",
- stochastic_weight_avg: bool = False,
- terminate_on_nan: Optional[bool] = None,
- ) -> None:
- r"""
- Customize every aspect of training via flags.
-
- Args:
-
- accelerator: Supports passing different accelerator types ("cpu", "gpu", "tpu", "ipu", "hpu", "auto")
- as well as custom accelerator instances.
-
- .. deprecated:: v1.5
- Passing training strategies (e.g., 'ddp') to ``accelerator`` has been deprecated in v1.5.0
- and will be removed in v1.7.0. Please use the ``strategy`` argument instead.
-
- accumulate_grad_batches: Accumulates grads every k batches or as set up in the dict.
- Default: ``None``.
-
- amp_backend: The mixed precision backend to use ("native" or "apex").
- Default: ``'native''``.
-
- amp_level: The optimization level to use (O1, O2, etc...). By default it will be set to "O2"
- if ``amp_backend`` is set to "apex".
-
- auto_lr_find: If set to True, will make trainer.tune() run a learning rate finder,
- trying to optimize initial learning for faster convergence. trainer.tune() method will
- set the suggested learning rate in self.lr or self.learning_rate in the LightningModule.
- To use a different key set a string instead of True with the key name.
- Default: ``False``.
-
- auto_scale_batch_size: If set to True, will `initially` run a batch size
- finder trying to find the largest batch size that fits into memory.
- The result will be stored in self.batch_size in the LightningModule.
- Additionally, can be set to either `power` that estimates the batch size through
- a power search or `binsearch` that estimates the batch size through a binary search.
- Default: ``False``.
-
- auto_select_gpus: If enabled and ``gpus`` or ``devices`` is an integer, pick available
- gpus automatically. This is especially useful when
- GPUs are configured to be in "exclusive mode", such
- that only one process at a time can access them.
- Default: ``False``.
-
- benchmark: Sets ``torch.backends.cudnn.benchmark``.
- Defaults to ``True`` if :paramref:`~pytorch_lightning.trainer.trainer.Trainer.deterministic`
- is ``False``. Overwrite to manually set a different value. Default: ``None``.
-
- callbacks: Add a callback or list of callbacks.
- Default: ``None``.
-
- checkpoint_callback: If ``True``, enable checkpointing.
- Default: ``None``.
-
- .. deprecated:: v1.5
- ``checkpoint_callback`` has been deprecated in v1.5 and will be removed in v1.7.
- Please consider using ``enable_checkpointing`` instead.
-
- enable_checkpointing: If ``True``, enable checkpointing.
- It will configure a default ModelCheckpoint callback if there is no user-defined ModelCheckpoint in
- :paramref:`~pytorch_lightning.trainer.trainer.Trainer.callbacks`.
- Default: ``True``.
-
- check_val_every_n_epoch: Check val every n train epochs.
- Default: ``1``.
-
-
- default_root_dir: Default path for logs and weights when no logger/ckpt_callback passed.
- Default: ``os.getcwd()``.
- Can be remote file paths such as `s3://mybucket/path` or 'hdfs://path/'
-
- detect_anomaly: Enable anomaly detection for the autograd engine.
- Default: ``False``.
-
- deterministic: If ``True``, sets whether PyTorch operations must use deterministic algorithms.
- Default: ``False``.
-
- devices: Will be mapped to either `gpus`, `tpu_cores`, `num_processes` or `ipus`,
- based on the accelerator type.
-
- fast_dev_run: Runs n if set to ``n`` (int) else 1 if set to ``True`` batch(es)
- of train, val and test to find any bugs (ie: a sort of unit test).
- Default: ``False``.
-
- flush_logs_every_n_steps: How often to flush logs to disk (defaults to every 100 steps).
-
- .. deprecated:: v1.5
- ``flush_logs_every_n_steps`` has been deprecated in v1.5 and will be removed in v1.7.
- Please configure flushing directly in the logger instead.
-
- gpus: Number of GPUs to train on (int) or which GPUs to train on (list or str) applied per node
- Default: ``None``.
-
- gradient_clip_val: The value at which to clip gradients. Passing ``gradient_clip_val=None`` disables
- gradient clipping. If using Automatic Mixed Precision (AMP), the gradients will be unscaled before.
- Default: ``None``.
-
- gradient_clip_algorithm: The gradient clipping algorithm to use. Pass ``gradient_clip_algorithm="value"``
- to clip by value, and ``gradient_clip_algorithm="norm"`` to clip by norm. By default it will
- be set to ``"norm"``.
-
- limit_train_batches: How much of training dataset to check (float = fraction, int = num_batches).
- Default: ``1.0``.
-
- limit_val_batches: How much of validation dataset to check (float = fraction, int = num_batches).
- Default: ``1.0``.
-
- limit_test_batches: How much of test dataset to check (float = fraction, int = num_batches).
- Default: ``1.0``.
-
- limit_predict_batches: How much of prediction dataset to check (float = fraction, int = num_batches).
- Default: ``1.0``.
-
- logger: Logger (or iterable collection of loggers) for experiment tracking. A ``True`` value uses
- the default ``TensorBoardLogger``. ``False`` will disable logging. If multiple loggers are
- provided and the `save_dir` property of that logger is not set, local files (checkpoints,
- profiler traces, etc.) are saved in ``default_root_dir`` rather than in the ``log_dir`` of any
- of the individual loggers.
- Default: ``True``.
-
- log_gpu_memory: None, 'min_max', 'all'. Might slow performance.
-
- .. deprecated:: v1.5
- Deprecated in v1.5.0 and will be removed in v1.7.0
- Please use the ``DeviceStatsMonitor`` callback directly instead.
-
- log_every_n_steps: How often to log within steps.
- Default: ``50``.
-
- prepare_data_per_node: If True, each LOCAL_RANK=0 will call prepare data.
- Otherwise only NODE_RANK=0, LOCAL_RANK=0 will prepare data
-
- .. deprecated:: v1.5
- Deprecated in v1.5.0 and will be removed in v1.7.0
- Please set ``prepare_data_per_node`` in ``LightningDataModule`` and/or
- ``LightningModule`` directly instead.
-
- process_position: Orders the progress bar when running multiple models on same machine.
-
- .. deprecated:: v1.5
- ``process_position`` has been deprecated in v1.5 and will be removed in v1.7.
- Please pass :class:`~pytorch_lightning.callbacks.progress.TQDMProgressBar` with ``process_position``
- directly to the Trainer's ``callbacks`` argument instead.
-
- progress_bar_refresh_rate: How often to refresh progress bar (in steps). Value ``0`` disables progress bar.
- Ignored when a custom progress bar is passed to :paramref:`~Trainer.callbacks`. Default: None, means
- a suitable value will be chosen based on the environment (terminal, Google COLAB, etc.).
-
- .. deprecated:: v1.5
- ``progress_bar_refresh_rate`` has been deprecated in v1.5 and will be removed in v1.7.
- Please pass :class:`~pytorch_lightning.callbacks.progress.TQDMProgressBar` with ``refresh_rate``
- directly to the Trainer's ``callbacks`` argument instead. To disable the progress bar,
- pass ``enable_progress_bar = False`` to the Trainer.
-
- enable_progress_bar: Whether to enable to progress bar by default.
- Default: ``False``.
-
- profiler: To profile individual steps during training and assist in identifying bottlenecks.
- Default: ``None``.
-
- overfit_batches: Overfit a fraction of training data (float) or a set number of batches (int).
- Default: ``0.0``.
-
- plugins: Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins.
- Default: ``None``.
-
- precision: Double precision (64), full precision (32), half precision (16) or bfloat16 precision (bf16).
- Can be used on CPU, GPU, TPUs, HPUs or IPUs.
- Default: ``32``.
-
- max_epochs: Stop training once this number of epochs is reached. Disabled by default (None).
- If both max_epochs and max_steps are not specified, defaults to ``max_epochs = 1000``.
- To enable infinite training, set ``max_epochs = -1``.
-
- min_epochs: Force training for at least these many epochs. Disabled by default (None).
-
- max_steps: Stop training after this number of steps. Disabled by default (-1). If ``max_steps = -1``
- and ``max_epochs = None``, will default to ``max_epochs = 1000``. To enable infinite training, set
- ``max_epochs`` to ``-1``.
-
- min_steps: Force training for at least these number of steps. Disabled by default (``None``).
-
- max_time: Stop training after this amount of time has passed. Disabled by default (``None``).
- The time duration can be specified in the format DD:HH:MM:SS (days, hours, minutes seconds), as a
- :class:`datetime.timedelta`, or a dictionary with keys that will be passed to
- :class:`datetime.timedelta`.
-
- num_nodes: Number of GPU nodes for distributed training.
- Default: ``1``.
-
- num_processes: Number of processes for distributed training with ``accelerator="cpu"``.
- Default: ``1``.
-
- num_sanity_val_steps: Sanity check runs n validation batches before starting the training routine.
- Set it to `-1` to run all batches in all validation dataloaders.
- Default: ``2``.
-
- reload_dataloaders_every_n_epochs: Set to a non-negative integer to reload dataloaders every n epochs.
- Default: ``0``.
-
- replace_sampler_ddp: Explicitly enables or disables sampler replacement. If not specified this
- will toggled automatically when DDP is used. By default it will add ``shuffle=True`` for
- train sampler and ``shuffle=False`` for val/test sampler. If you want to customize it,
- you can set ``replace_sampler_ddp=False`` and add your own distributed sampler.
-
- resume_from_checkpoint: Path/URL of the checkpoint from which training is resumed. If there is
- no checkpoint file at the path, an exception is raised. If resuming from mid-epoch checkpoint,
- training will start from the beginning of the next epoch.
-
- .. deprecated:: v1.5
- ``resume_from_checkpoint`` is deprecated in v1.5 and will be removed in v2.0.
- Please pass the path to ``Trainer.fit(..., ckpt_path=...)`` instead.
-
- strategy: Supports different training strategies with aliases
- as well custom strategies.
- Default: ``None``.
-
- sync_batchnorm: Synchronize batch norm layers between process groups/whole world.
- Default: ``False``.
-
- terminate_on_nan: If set to True, will terminate training (by raising a `ValueError`) at the
- end of each training batch, if any of the parameters or the loss are NaN or +/-inf.
-
- .. deprecated:: v1.5
- Trainer argument ``terminate_on_nan`` was deprecated in v1.5 and will be removed in 1.7.
- Please use ``detect_anomaly`` instead.
-
- detect_anomaly: Enable anomaly detection for the autograd engine.
- Default: ``False``.
-
- tpu_cores: How many TPU cores to train on (1 or 8) / Single TPU to train on (1)
- Default: ``None``.
-
- ipus: How many IPUs to train on.
- Default: ``None``.
-
- track_grad_norm: -1 no tracking. Otherwise tracks that p-norm. May be set to 'inf' infinity-norm. If using
- Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
- Default: ``-1``.
-
- val_check_interval: How often to check the validation set. Pass a ``float`` in the range [0.0, 1.0] to check
- after a fraction of the training epoch. Pass an ``int`` to check after a fixed number of training
- batches.
- Default: ``1.0``.
-
- enable_model_summary: Whether to enable model summarization by default.
- Default: ``True``.
-
- weights_summary: Prints a summary of the weights when training begins.
-
- .. deprecated:: v1.5
- ``weights_summary`` has been deprecated in v1.5 and will be removed in v1.7.
- To disable the summary, pass ``enable_model_summary = False`` to the Trainer.
- To customize the summary, pass :class:`~pytorch_lightning.callbacks.model_summary.ModelSummary`
- directly to the Trainer's ``callbacks`` argument.
-
- weights_save_path: Where to save weights if specified. Will override default_root_dir
- for checkpoints only. Use this if for whatever reason you need the checkpoints
- stored in a different place than the logs written in `default_root_dir`.
- Can be remote file paths such as `s3://mybucket/path` or 'hdfs://path/'
- Defaults to `default_root_dir`.
-
- .. deprecated:: v1.6
- ``weights_save_path`` has been deprecated in v1.6 and will be removed in v1.8. Please pass
- ``dirpath`` directly to the :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint`
- callback.
-
- move_metrics_to_cpu: Whether to force internal logged metrics to be moved to cpu.
- This can save some gpu memory, but can make training slower. Use with attention.
- Default: ``False``.
-
- multiple_trainloader_mode: How to loop over the datasets when there are multiple train loaders.
- In 'max_size_cycle' mode, the trainer ends one epoch when the largest dataset is traversed,
- and smaller datasets reload when running out of their data. In 'min_size' mode, all the datasets
- reload when reaching the minimum length of datasets.
- Default: ``"max_size_cycle"``.
-
- stochastic_weight_avg: Whether to use `Stochastic Weight Averaging (SWA)
- <https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/>`_.
- Default: ``False``.
-
- .. deprecated:: v1.5
- ``stochastic_weight_avg`` has been deprecated in v1.5 and will be removed in v1.7.
- Please pass :class:`~pytorch_lightning.callbacks.stochastic_weight_avg.StochasticWeightAveraging`
- directly to the Trainer's ``callbacks`` argument instead.
- """
- super().__init__()
- Trainer._log_api_event("init")
- log.detail(f"{self.__class__.__name__}: Initializing trainer with parameters: {locals()}")
- self.state = TrainerState()
-
- # init connectors
- self._data_connector = DataConnector(self, multiple_trainloader_mode)
-
- self._accelerator_connector = AcceleratorConnector(
- num_processes=num_processes,
- devices=devices,
- tpu_cores=tpu_cores,
- ipus=ipus,
- accelerator=accelerator,
- strategy=strategy,
- gpus=gpus,
- num_nodes=num_nodes,
- sync_batchnorm=sync_batchnorm,
- benchmark=benchmark,
- replace_sampler_ddp=replace_sampler_ddp,
- deterministic=deterministic,
- auto_select_gpus=auto_select_gpus,
- precision=precision,
- amp_type=amp_backend,
- amp_level=amp_level,
- plugins=plugins,
- )
- self._logger_connector = LoggerConnector(self, log_gpu_memory)
- self._callback_connector = CallbackConnector(self)
- self._checkpoint_connector = CheckpointConnector(self, resume_from_checkpoint)
- self._signal_connector = SignalConnector(self)
- self.tuner = Tuner(self)
-
- min_steps, max_steps, min_epochs, max_epochs, max_time = _parse_loop_limits(
- min_steps, max_steps, min_epochs, max_epochs, max_time
- )
- fit_loop = FitLoop(min_epochs=min_epochs, max_epochs=max_epochs)
- training_epoch_loop = TrainingEpochLoop(min_steps=min_steps, max_steps=max_steps)
- fit_loop.connect(epoch_loop=training_epoch_loop)
-
- # default .fit() loop
- self.fit_loop = fit_loop
-
- # default .validate() loop
- self.validate_loop = EvaluationLoop()
-
- # default .test() loop
- self.test_loop = EvaluationLoop()
-
- # default .predict() loop
- self.predict_loop = PredictionLoop()
-
- # set when a checkpoint is loaded via `Trainer.{fit,validate,test,predict}`.
- self._ckpt_path: Optional[str] = None
-
- # .validate(), predict() and .test() set these when they load a checkpoint. They will be removed in favor of
- # the unified read-only `Trainer.ckpt_path` attribute in v1.8
- self._validated_ckpt_path: Optional[str] = None # TODO: remove in v1.8
- self._tested_ckpt_path: Optional[str] = None # TODO: remove in v1.8
- self._predicted_ckpt_path: Optional[str] = None # TODO: remove in v1.8
-
- # todo: remove in v1.7
- self._weights_summary: Optional[str] = None
-
- # init callbacks
- # Declare attributes to be set in _callback_connector on_trainer_init
- self._callback_connector.on_trainer_init(
- callbacks,
- checkpoint_callback,
- enable_checkpointing,
- enable_progress_bar,
- progress_bar_refresh_rate,
- process_position,
- default_root_dir,
- weights_save_path,
- enable_model_summary,
- weights_summary,
- stochastic_weight_avg,
- max_time,
- accumulate_grad_batches,
- )
-
- # hook
- self._call_callback_hooks("on_init_start")
-
- # init data flags
- self.check_val_every_n_epoch: int
- self._data_connector.on_trainer_init(
- check_val_every_n_epoch,
- reload_dataloaders_every_n_epochs,
- prepare_data_per_node,
- )
-
- if terminate_on_nan is not None:
- rank_zero_deprecation(
- "Trainer argument `terminate_on_nan` was deprecated in v1.5 and will be removed in 1.7."
- " Please use `Trainer(detect_anomaly=True)` instead."
- )
- if not isinstance(terminate_on_nan, bool):
- raise TypeError(f"`terminate_on_nan` should be a bool, got {terminate_on_nan}.")
-
- # gradient clipping
- if gradient_clip_val is not None and not isinstance(gradient_clip_val, (int, float)):
- raise TypeError(f"`gradient_clip_val` should be an int or a float. Got {gradient_clip_val}.")
-
- if gradient_clip_algorithm is not None and not GradClipAlgorithmType.supported_type(
- gradient_clip_algorithm.lower()
- ):
- raise MisconfigurationException(
- f"`gradient_clip_algorithm` {gradient_clip_algorithm} is invalid. "
- f"Allowed algorithms: {GradClipAlgorithmType.supported_types()}."
- )
-
- # gradient norm tracking
- if track_grad_norm != -1 and not (
- (isinstance(track_grad_norm, (int, float)) or track_grad_norm == "inf") and float(track_grad_norm) > 0
- ):
- raise MisconfigurationException(
- f"`track_grad_norm` must be a positive number or 'inf' (infinity norm). Got {track_grad_norm}."
- )
-
- self._terminate_on_nan = terminate_on_nan
- self.gradient_clip_val: Union[int, float] = gradient_clip_val
- self.gradient_clip_algorithm: Optional[GradClipAlgorithmType] = (
- GradClipAlgorithmType(gradient_clip_algorithm.lower()) if gradient_clip_algorithm is not None else None
- )
- self.track_grad_norm: float = float(track_grad_norm)
-
- self._detect_anomaly: bool = detect_anomaly
- self._setup_on_init(num_sanity_val_steps)
-
- # configure tuner
- self.tuner.on_trainer_init(auto_lr_find, auto_scale_batch_size)
-
- # configure profiler
- self.__init_profiler(profiler)
-
- # init logger flags
- self._loggers: List[LightningLoggerBase]
- self._logger_connector.on_trainer_init(logger, flush_logs_every_n_steps, log_every_n_steps, move_metrics_to_cpu)
-
- # init debugging flags
- self.val_check_interval: Union[int, float]
- self._init_debugging_flags(
- limit_train_batches,
- limit_val_batches,
- limit_test_batches,
- limit_predict_batches,
- val_check_interval,
- overfit_batches,
- fast_dev_run,
- )
-
- # Callback system
- self._call_callback_hooks("on_init_end")
-
- def _init_debugging_flags(
- self,
- limit_train_batches: Optional[Union[int, float]],
- limit_val_batches: Optional[Union[int, float]],
- limit_test_batches: Optional[Union[int, float]],
- limit_predict_batches: Optional[Union[int, float]],
- val_check_interval: Optional[Union[int, float]],
- overfit_batches: Union[int, float],
- fast_dev_run: Union[int, bool],
- ) -> None:
- if isinstance(fast_dev_run, int) and (fast_dev_run < 0):
- raise MisconfigurationException(
- f"fast_dev_run={fast_dev_run} is not a valid configuration. It should be >= 0."
- )
-
- self.fast_dev_run = fast_dev_run
-
- # set fast_dev_run=True when it is 1, used while logging
- if fast_dev_run == 1:
- self.fast_dev_run = True
-
- if fast_dev_run:
- num_batches = int(fast_dev_run)
- limit_train_batches = num_batches
- limit_val_batches = num_batches
- limit_test_batches = num_batches
- limit_predict_batches = num_batches
- self.fit_loop.max_steps = num_batches
- self.num_sanity_val_steps = 0
- self.fit_loop.max_epochs = 1
- val_check_interval = 1.0
- self.check_val_every_n_epoch = 1
- self.loggers = [DummyLogger()] if self.loggers else []
-
- rank_zero_info(
- "Running in fast_dev_run mode: will run a full train,"
- f" val, test and prediction loop using {num_batches} batch(es)."
- )
-
- self.limit_train_batches = _determine_batch_limits(limit_train_batches, "limit_train_batches")
- self.limit_val_batches = _determine_batch_limits(limit_val_batches, "limit_val_batches")
- self.limit_test_batches = _determine_batch_limits(limit_test_batches, "limit_test_batches")
- self.limit_predict_batches = _determine_batch_limits(limit_predict_batches, "limit_predict_batches")
- self.val_check_interval = _determine_batch_limits(val_check_interval, "val_check_interval")
- self.overfit_batches = _determine_batch_limits(overfit_batches, "overfit_batches")
- self._determine_data_use_amount(self.overfit_batches)
-
- def _determine_data_use_amount(self, overfit_batches: float) -> None:
- """Use less data for debugging purposes."""
- if overfit_batches > 0:
- self.limit_train_batches = overfit_batches
- self.limit_val_batches = 0
-
- def _setup_on_init(self, num_sanity_val_steps: int) -> None:
- self._log_device_info()
-
- self.should_stop = False
- self.state = TrainerState()
- self.num_training_batches = float("inf")
- self.train_dataloader = None
-
- if num_sanity_val_steps == -1:
- self.num_sanity_val_steps = float("inf")
- else:
- self.num_sanity_val_steps = num_sanity_val_steps
-
- self.num_sanity_val_batches = []
- self.num_test_batches = []
- self.num_val_batches = []
- self.test_dataloaders = None
- self.val_dataloaders = None
- self._last_train_dl_reload_epoch = float("-inf")
- self._last_val_dl_reload_epoch = float("-inf")
-
- self.num_predict_batches = []
-
- def _call_and_handle_interrupt(self, trainer_fn: Callable, *args: Any, **kwargs: Any) -> Any:
- r"""
- Error handling, intended to be used only for main trainer function entry points (fit, validate, test, predict)
- as all errors should funnel through them
-
- Args:
- trainer_fn: one of (fit, validate, test, predict)
- *args: positional arguments to be passed to the `trainer_fn`
- **kwargs: keyword arguments to be passed to `trainer_fn`
- """
- try:
- if self.strategy.launcher is not None:
- return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
- else:
- return trainer_fn(*args, **kwargs)
- # TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7
- except KeyboardInterrupt as exception:
- rank_zero_warn("Detected KeyboardInterrupt, attempting graceful shutdown...")
- # user could press Ctrl+c many times... only shutdown once
- if not self.interrupted:
- self.state.status = TrainerStatus.INTERRUPTED
- self._call_callback_hooks("on_keyboard_interrupt")
- self._call_callback_hooks("on_exception", exception)
- except BaseException as exception:
- self.state.status = TrainerStatus.INTERRUPTED
- if distributed_available() and self.world_size > 1:
- # try syncing remaining processes, kill otherwise
- self.strategy.reconciliate_processes(traceback.format_exc())
- self._call_callback_hooks("on_exception", exception)
- self._teardown()
- # teardown might access the stage so we reset it after
- self.state.stage = None
- raise
-
-[docs] def fit(
- self,
- model: "pl.LightningModule",
- train_dataloaders: Optional[Union[TRAIN_DATALOADERS, LightningDataModule]] = None,
- val_dataloaders: Optional[EVAL_DATALOADERS] = None,
- datamodule: Optional[LightningDataModule] = None,
- ckpt_path: Optional[str] = None,
- ) -> None:
- r"""
- Runs the full optimization routine.
-
- Args:
- model: Model to fit.
-
- train_dataloaders: A collection of :class:`torch.utils.data.DataLoader` or a
- :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying training samples.
- In the case of multiple dataloaders, please see this :ref:`section <multiple-dataloaders>`.
-
- val_dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them specifying validation samples.
-
- ckpt_path: Path/URL of the checkpoint from which training is resumed. If there is
- no checkpoint file at the path, an exception is raised. If resuming from mid-epoch checkpoint,
- training will start from the beginning of the next epoch.
-
- datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
- """
- self.strategy.model = model
- self._call_and_handle_interrupt(
- self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
- )
-
- def _fit_impl(
- self,
- model: "pl.LightningModule",
- train_dataloaders: Optional[Union[TRAIN_DATALOADERS, LightningDataModule]] = None,
- val_dataloaders: Optional[EVAL_DATALOADERS] = None,
- datamodule: Optional[LightningDataModule] = None,
- ckpt_path: Optional[str] = None,
- ) -> None:
- Trainer._log_api_event("fit")
- log.detail(f"{self.__class__.__name__}: trainer fit stage")
-
- self.state.fn = TrainerFn.FITTING
- self.state.status = TrainerStatus.RUNNING
- self.training = True
- self._last_train_dl_reload_epoch = float("-inf")
- self._last_val_dl_reload_epoch = float("-inf")
-
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(train_dataloaders, LightningDataModule):
- datamodule = train_dataloaders
- train_dataloaders = None
- # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
- if (train_dataloaders is not None or val_dataloaders is not None) and datamodule is not None:
- raise MisconfigurationException(
- "You cannot pass `train_dataloader` or `val_dataloaders` to `trainer.fit(datamodule=...)`"
- )
-
- # links data to the trainer
- self._data_connector.attach_data(
- model, train_dataloaders=train_dataloaders, val_dataloaders=val_dataloaders, datamodule=datamodule
- )
-
- # TODO: ckpt_path only in v2.0
- ckpt_path = ckpt_path or self.resume_from_checkpoint
- self._ckpt_path = self.__set_ckpt_path(
- ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
- )
- results = self._run(model, ckpt_path=self.ckpt_path)
-
- assert self.state.stopped
- self.training = False
- return results
-
-[docs] def validate(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- ckpt_path: Optional[str] = None,
- verbose: bool = True,
- datamodule: Optional[LightningDataModule] = None,
- ) -> _EVALUATE_OUTPUT:
- r"""
- Perform one evaluation epoch over the validation set.
-
- Args:
- model: The model to validate.
-
- dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them,
- or a :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying validation samples.
-
- ckpt_path: Either ``best`` or path to the checkpoint you wish to validate.
- If ``None`` and the model instance was passed, use the current weights.
- Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
- if a checkpoint callback is configured.
-
- verbose: If True, prints the validation results.
-
- datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
-
- Returns:
- List of dictionaries with metrics logged during the validation phase, e.g., in model- or callback hooks
- like :meth:`~pytorch_lightning.core.lightning.LightningModule.validation_step`,
- :meth:`~pytorch_lightning.core.lightning.LightningModule.validation_epoch_end`, etc.
- The length of the list corresponds to the number of validation dataloaders used.
- """
- self.strategy.model = model or self.lightning_module
- return self._call_and_handle_interrupt(self._validate_impl, model, dataloaders, ckpt_path, verbose, datamodule)
-
- def _validate_impl(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- ckpt_path: Optional[str] = None,
- verbose: bool = True,
- datamodule: Optional[LightningDataModule] = None,
- ) -> _EVALUATE_OUTPUT:
- # --------------------
- # SETUP HOOK
- # --------------------
- Trainer._log_api_event("validate")
- log.detail(f"{self.__class__.__name__}: trainer validate stage")
-
- self.state.fn = TrainerFn.VALIDATING
- self.state.status = TrainerStatus.RUNNING
- self.validating = True
-
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(dataloaders, LightningDataModule):
- datamodule = dataloaders
- dataloaders = None
- # If you supply a datamodule you can't supply val_dataloaders
- if dataloaders is not None and datamodule:
- raise MisconfigurationException("You cannot pass both `trainer.validate(dataloaders=..., datamodule=...)`")
-
- model_provided = model is not None
- model = model or self.lightning_module
- if model is None:
- raise MisconfigurationException(
- "`model` must be provided to `trainer.validate()` when it hasn't been passed in a previous run"
- )
-
- self.validate_loop.verbose = verbose
-
- # links data to the trainer
- self._data_connector.attach_data(model, val_dataloaders=dataloaders, datamodule=datamodule)
-
- self._ckpt_path = self.__set_ckpt_path(
- ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
- )
-
- self._validated_ckpt_path = self.ckpt_path # TODO: remove in v1.8
-
- # run validate
- results = self._run(model, ckpt_path=self.ckpt_path)
-
- assert self.state.stopped
- self.validating = False
-
- return results
-
-[docs] def test(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- ckpt_path: Optional[str] = None,
- verbose: bool = True,
- datamodule: Optional[LightningDataModule] = None,
- ) -> _EVALUATE_OUTPUT:
- r"""
- Perform one evaluation epoch over the test set.
- It's separated from fit to make sure you never run on your test set until you want to.
-
- Args:
- model: The model to test.
-
- dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them,
- or a :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying test samples.
-
- ckpt_path: Either ``best`` or path to the checkpoint you wish to test.
- If ``None`` and the model instance was passed, use the current weights.
- Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
- if a checkpoint callback is configured.
-
- verbose: If True, prints the test results.
-
- datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
-
- Returns:
- List of dictionaries with metrics logged during the test phase, e.g., in model- or callback hooks
- like :meth:`~pytorch_lightning.core.lightning.LightningModule.test_step`,
- :meth:`~pytorch_lightning.core.lightning.LightningModule.test_epoch_end`, etc.
- The length of the list corresponds to the number of test dataloaders used.
- """
- self.strategy.model = model or self.lightning_module
- return self._call_and_handle_interrupt(self._test_impl, model, dataloaders, ckpt_path, verbose, datamodule)
-
- def _test_impl(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- ckpt_path: Optional[str] = None,
- verbose: bool = True,
- datamodule: Optional[LightningDataModule] = None,
- ) -> _EVALUATE_OUTPUT:
- # --------------------
- # SETUP HOOK
- # --------------------
- Trainer._log_api_event("test")
- log.detail(f"{self.__class__.__name__}: trainer test stage")
-
- self.state.fn = TrainerFn.TESTING
- self.state.status = TrainerStatus.RUNNING
- self.testing = True
-
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(dataloaders, LightningDataModule):
- datamodule = dataloaders
- dataloaders = None
- # If you supply a datamodule you can't supply test_dataloaders
- if dataloaders is not None and datamodule:
- raise MisconfigurationException("You cannot pass both `trainer.test(dataloaders=..., datamodule=...)`")
-
- model_provided = model is not None
- model = model or self.lightning_module
- if model is None:
- raise MisconfigurationException(
- "`model` must be provided to `trainer.test()` when it hasn't been passed in a previous run"
- )
-
- self.test_loop.verbose = verbose
-
- # links data to the trainer
- self._data_connector.attach_data(model, test_dataloaders=dataloaders, datamodule=datamodule)
-
- self._ckpt_path = self.__set_ckpt_path(
- ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
- )
-
- self._tested_ckpt_path = self.ckpt_path # TODO: remove in v1.8
-
- # run test
- results = self._run(model, ckpt_path=self.ckpt_path)
-
- assert self.state.stopped
- self.testing = False
-
- return results
-
-[docs] def predict(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- datamodule: Optional[LightningDataModule] = None,
- return_predictions: Optional[bool] = None,
- ckpt_path: Optional[str] = None,
- ) -> Optional[_PREDICT_OUTPUT]:
- r"""
- Run inference on your data.
- This will call the model forward function to compute predictions. Useful to perform distributed
- and batched predictions. Logging is disabled in the predict hooks.
-
- Args:
- model: The model to predict with.
-
- dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them,
- or a :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying prediction samples.
-
- datamodule: The datamodule with a predict_dataloader method that returns one or more dataloaders.
-
- return_predictions: Whether to return predictions.
- ``True`` by default except when an accelerator that spawns processes is used (not supported).
-
- ckpt_path: Either ``best`` or path to the checkpoint you wish to predict.
- If ``None`` and the model instance was passed, use the current weights.
- Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
- if a checkpoint callback is configured.
-
- Returns:
- Returns a list of dictionaries, one for each provided dataloader containing their respective predictions.
- """
- self.strategy.model = model or self.lightning_module
- return self._call_and_handle_interrupt(
- self._predict_impl, model, dataloaders, datamodule, return_predictions, ckpt_path
- )
-
- def _predict_impl(
- self,
- model: Optional["pl.LightningModule"] = None,
- dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
- datamodule: Optional[LightningDataModule] = None,
- return_predictions: Optional[bool] = None,
- ckpt_path: Optional[str] = None,
- ) -> Optional[_PREDICT_OUTPUT]:
- # --------------------
- # SETUP HOOK
- # --------------------
- Trainer._log_api_event("predict")
- log.detail(f"{self.__class__.__name__}: trainer predict stage")
-
- self.state.fn = TrainerFn.PREDICTING
- self.state.status = TrainerStatus.RUNNING
- self.predicting = True
-
- self.predict_loop.return_predictions = return_predictions
-
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(dataloaders, LightningDataModule):
- datamodule = dataloaders
- dataloaders = None
- if dataloaders is not None and datamodule:
- raise MisconfigurationException("You cannot pass both `trainer.predict(dataloaders=..., datamodule=...)`")
-
- model_provided = model is not None
- model = model or self.lightning_module
- if model is None:
- raise MisconfigurationException(
- "`model` must be provided to `trainer.predict()` when it hasn't been passed in a previous run"
- )
-
- # links data to the trainer
- self._data_connector.attach_data(model, predict_dataloaders=dataloaders, datamodule=datamodule)
-
- self._ckpt_path = self.__set_ckpt_path(
- ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
- )
-
- self._predicted_ckpt_path = self.ckpt_path # TODO: remove in v1.8
-
- results = self._run(model, ckpt_path=self.ckpt_path)
-
- assert self.state.stopped
- self.predicting = False
-
- return results
-
-[docs] def tune(
- self,
- model: "pl.LightningModule",
- train_dataloaders: Optional[Union[TRAIN_DATALOADERS, LightningDataModule]] = None,
- val_dataloaders: Optional[EVAL_DATALOADERS] = None,
- datamodule: Optional[LightningDataModule] = None,
- scale_batch_size_kwargs: Optional[Dict[str, Any]] = None,
- lr_find_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Dict[str, Optional[Union[int, _LRFinder]]]:
- r"""
- Runs routines to tune hyperparameters before training.
-
- Args:
- model: Model to tune.
-
- train_dataloaders: A collection of :class:`torch.utils.data.DataLoader` or a
- :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying training samples.
- In the case of multiple dataloaders, please see this :ref:`section <multiple-dataloaders>`.
-
- val_dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them specifying validation samples.
-
- datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
-
- scale_batch_size_kwargs: Arguments for :func:`~pytorch_lightning.tuner.batch_size_scaling.scale_batch_size`
-
- lr_find_kwargs: Arguments for :func:`~pytorch_lightning.tuner.lr_finder.lr_find`
- """
- Trainer._log_api_event("tune")
-
- self.state.fn = TrainerFn.TUNING
- self.state.status = TrainerStatus.RUNNING
- self.tuning = True
-
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(train_dataloaders, LightningDataModule):
- datamodule = train_dataloaders
- train_dataloaders = None
- # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
- if (train_dataloaders is not None or val_dataloaders is not None) and datamodule is not None:
- raise MisconfigurationException(
- "You cannot pass `train_dataloader` or `val_dataloaders` to `trainer.tune(datamodule=...)`"
- )
-
- # links data to the trainer
- self._data_connector.attach_data(
- model, train_dataloaders=train_dataloaders, val_dataloaders=val_dataloaders, datamodule=datamodule
- )
-
- with isolate_rng():
- result = self.tuner._tune(
- model, scale_batch_size_kwargs=scale_batch_size_kwargs, lr_find_kwargs=lr_find_kwargs
- )
-
- assert self.state.stopped
- self.tuning = False
-
- return result
-
- def _restore_modules_and_callbacks(self, checkpoint_path: Optional[_PATH] = None) -> None:
- # restore modules after setup
- self._checkpoint_connector.resume_start(checkpoint_path)
- self._checkpoint_connector._restore_quantization_callbacks()
- self._checkpoint_connector.restore_model()
- self._checkpoint_connector.restore_datamodule()
- if self.state.fn == TrainerFn.FITTING:
- # restore callback states
- self._checkpoint_connector.restore_callbacks()
-
- def _run(
- self, model: "pl.LightningModule", ckpt_path: Optional[str] = None
- ) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]:
- # clean hparams
- if hasattr(model, "hparams"):
- parsing.clean_namespace(model.hparams)
-
- # attach model to the strategy
- self.strategy.connect(model)
-
- self._callback_connector._attach_model_callbacks()
- self._callback_connector._attach_model_logging_functions()
-
- verify_loop_configurations(self)
-
- # hook
- log.detail(f"{self.__class__.__name__}: preparing data")
- self._data_connector.prepare_data()
-
- # ----------------------------
- # SET UP TRAINING
- # ----------------------------
- self._call_callback_hooks("on_before_accelerator_backend_setup")
- log.detail(f"{self.__class__.__name__}: setting up strategy environment")
- self.strategy.setup_environment()
- self.__setup_profiler()
-
- self._call_setup_hook() # allow user to setup lightning_module in accelerator environment
-
- # check if we should delay restoring checkpoint till later
- if not self.strategy.restore_checkpoint_after_setup:
- log.detail(f"{self.__class__.__name__}: restoring module and callbacks from checkpoint path: {ckpt_path}")
- self._restore_modules_and_callbacks(ckpt_path)
-
- log.detail(f"{self.__class__.__name__}: configuring sharded model")
- self._call_configure_sharded_model() # allow user to setup in model sharded environment
-
- # ----------------------------
- # INSPECT THE CORE LOOPS
- # ----------------------------
- rf"""
- Lightning internal flow looks like this:
- {Trainer.fit} or {Trainer.test} or {Trainer.predict} ||
- | ||
- spawn processes ||
- {self.strategy.setup_environment} ||
- | ||
- setup accelerator ||
- and strategy || LIGHTNING
- | ||
- {self._run_stage} || FLOW
- | ||
- {self._run_train} || DIRECTION
- or {self._run_evaluate} ||
- or {self._run_predict} ||
- | ||
- results \/
- This is used to guide readers to the core loops: train, test, predict.
- {self._run_predict} is the simplest to understand, use `Go to Definition` to read it :)
- """
-
- # ----------------------------
- # TRAIN
- # ----------------------------
-
- # reset logger connector
- self._logger_connector.reset_results()
- self._logger_connector.reset_metrics()
-
- # strategy will configure model and move it to the device
- self.strategy.setup(self)
-
- # hook
- if self.state.fn == TrainerFn.FITTING:
- self._call_callback_hooks("on_fit_start")
- self._call_lightning_module_hook("on_fit_start")
-
- self._log_hyperparams()
-
- if self.strategy.restore_checkpoint_after_setup:
- log.detail(f"{self.__class__.__name__}: restoring module and callbacks from checkpoint path: {ckpt_path}")
- self._restore_modules_and_callbacks(ckpt_path)
-
- # restore optimizers, etc.
- log.detail(f"{self.__class__.__name__}: restoring training state")
- self._checkpoint_connector.restore_training_state()
-
- self._checkpoint_connector.resume_end()
-
- results = self._run_stage()
-
- log.detail(f"{self.__class__.__name__}: trainer tearing down")
- self._teardown()
-
- # ----------------------------
- # POST-Training CLEAN UP
- # ----------------------------
- # hook
- if self.state.fn == TrainerFn.FITTING:
- self._call_callback_hooks("on_fit_end")
- self._call_lightning_module_hook("on_fit_end")
-
- log.detail(f"{self.__class__.__name__}: calling teardown hooks")
- self._call_teardown_hook()
-
- self.state.status = TrainerStatus.FINISHED
- self.state.stage = None
-
- return results
-
- def _log_hyperparams(self) -> None:
- if not self.loggers:
- return
- # log hyper-parameters
- hparams_initial = None
-
- # save exp to get started (this is where the first experiment logs are written)
- datamodule_log_hyperparams = self.datamodule._log_hyperparams if self.datamodule is not None else False
-
- if self.lightning_module._log_hyperparams and datamodule_log_hyperparams:
- datamodule_hparams = self.datamodule.hparams_initial
- lightning_hparams = self.lightning_module.hparams_initial
- inconsistent_keys = []
- for key in lightning_hparams.keys() & datamodule_hparams.keys():
- lm_val, dm_val = lightning_hparams[key], datamodule_hparams[key]
- if type(lm_val) != type(dm_val):
- inconsistent_keys.append(key)
- elif isinstance(lm_val, torch.Tensor) and id(lm_val) != id(dm_val):
- inconsistent_keys.append(key)
- elif lm_val != dm_val:
- inconsistent_keys.append(key)
- if inconsistent_keys:
- raise MisconfigurationException(
- f"Error while merging hparams: the keys {inconsistent_keys} are present "
- "in both the LightningModule's and LightningDataModule's hparams "
- "but have different values."
- )
- 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
-
- for logger in self.loggers:
- if hparams_initial is not None:
- logger.log_hyperparams(hparams_initial)
- logger.log_graph(self.lightning_module)
- logger.save()
-
- def _teardown(self):
- """This is the Trainer's internal teardown, unrelated to the `teardown` hooks in LightningModule and
- Callback; those are handled by :meth:`_call_teardown_hook`."""
- self.strategy.post_dispatch(self)
- self.strategy.teardown()
- loop = self._active_loop
- # loop should never be `None` here but it can because we don't know the trainer stage with `ddp_spawn`
- if loop is not None:
- loop.teardown()
- self._logger_connector.teardown()
- self._signal_connector.teardown()
-
- def run_stage(self) -> None:
- rank_zero_deprecation(
- "`Trainer.run_stage` is deprecated in v1.6 and will be removed in v1.8. Use"
- " `Trainer.{fit,validate,test,predict}` instead."
- )
- return self._run_stage()
-
- def _run_stage(self):
- self.strategy.barrier("run-stage")
- self.strategy.dispatch(self)
-
- if self.evaluating:
- return self._run_evaluate()
- if self.predicting:
- return self._run_predict()
- return self._run_train()
-
- def _pre_training_routine(self):
- # wait for all to join if on distributed
- self.strategy.barrier("setup_training")
-
- # register signals
- self._signal_connector.register_signal_handlers()
-
- # --------------------------
- # Pre-train
- # --------------------------
- self._call_callback_hooks("on_pretrain_routine_start")
- self._call_lightning_module_hook("on_pretrain_routine_start")
-
- self._call_callback_hooks("on_pretrain_routine_end")
- self._call_lightning_module_hook("on_pretrain_routine_end")
-
- def _run_train(self) -> None:
- self._pre_training_routine()
-
- with isolate_rng():
- self._run_sanity_check()
-
- # enable train mode
- self.model.train()
- torch.set_grad_enabled(True)
-
- self.fit_loop.trainer = self
- with torch.autograd.set_detect_anomaly(self._detect_anomaly):
- self.fit_loop.run()
-
- def _run_evaluate(self) -> _EVALUATE_OUTPUT:
- assert self.evaluating
-
- # reload dataloaders
- self._evaluation_loop._reload_evaluation_dataloaders()
-
- # reset trainer on this loop and all child loops in case user connected a custom loop
- self._evaluation_loop.trainer = self
-
- with self.profiler.profile(f"run_{self.state.stage}_evaluation"), torch.no_grad():
- eval_loop_results = self._evaluation_loop.run()
-
- # remove the tensors from the eval results
- for result in eval_loop_results:
- if isinstance(result, dict):
- for k, v in result.items():
- if isinstance(v, torch.Tensor):
- result[k] = v.cpu().item()
-
- return eval_loop_results
-
- def _run_predict(self) -> Optional[_PREDICT_OUTPUT]:
- self.reset_predict_dataloader(self.lightning_module)
- # reset trainer on this loop and all child loops in case user connected a custom loop
- self.predict_loop.trainer = self
- with torch.no_grad():
- return self.predict_loop.run()
-
- def _run_sanity_check(self) -> None:
- val_loop = self.fit_loop.epoch_loop.val_loop
-
- should_sanity_check = (
- self.enable_validation
- and self.num_sanity_val_steps > 0
- # do not sanity check if restarting because it would mess up the loaded state
- and not val_loop.restarting
- )
-
- # run tiny validation (if validation defined)
- # to make sure program won't crash during val
- if should_sanity_check:
- stage = self.state.stage
- self.sanity_checking = True
-
- # reset logger connector
- self._logger_connector.reset_results()
- self._logger_connector.reset_metrics()
-
- self._call_callback_hooks("on_sanity_check_start")
-
- # reload dataloaders
- val_loop._reload_evaluation_dataloaders()
- self.num_sanity_val_batches = [
- min(self.num_sanity_val_steps, val_batches) for val_batches in self.num_val_batches
- ]
-
- # run eval step
- with torch.no_grad():
- val_loop.run()
-
- self._call_callback_hooks("on_sanity_check_end")
-
- # reset logger connector
- self._logger_connector.reset_results()
- self._logger_connector.reset_metrics()
-
- # reset the progress tracking state after sanity checking. we don't need to set the state before
- # because sanity check only runs when we are not restarting
- _reset_progress(val_loop)
-
- # restore the previous stage when the sanity check if finished
- self.state.stage = stage
-
- def __set_ckpt_path(self, ckpt_path: Optional[str], model_provided: bool, model_connected: bool) -> Optional[str]:
- # fault-tolerance takes precedence
- from pytorch_lightning.callbacks.fault_tolerance import _FaultToleranceCheckpoint
-
- ft_checkpoints = [cb for cb in self.callbacks if isinstance(cb, _FaultToleranceCheckpoint)]
- if ft_checkpoints:
- ft_ckpt_path = ft_checkpoints[0].ckpt_path
- fs = get_filesystem(ft_ckpt_path)
- if fs.exists(ft_ckpt_path):
- return ft_ckpt_path
-
- if model_provided and ckpt_path is None:
- # use passed model to function without loading weights
- return
-
- fn = self.state.fn.value
-
- if model_connected and ckpt_path is None:
- rank_zero_warn(
- f"`.{fn}(ckpt_path=None)` was called without a model."
- " The best model of the previous `fit` call will be used."
- f" You can pass `{fn}(ckpt_path='best')` to use and best model"
- " checkpoint and avoid this warning or"
- " `ckpt_path=trainer.checkpoint_callback.last_model_path` to use the last model."
- )
- ckpt_path = "best"
-
- if ckpt_path == "best":
- if len(self.checkpoint_callbacks) > 1:
- rank_zero_warn(
- f'`.{fn}(ckpt_path="best")` is called with Trainer configured with multiple `ModelCheckpoint`'
- " callbacks. It will use the best checkpoint path from first checkpoint callback."
- )
-
- if not self.checkpoint_callback:
- raise MisconfigurationException(
- f'`.{fn}(ckpt_path="best")` is set but `ModelCheckpoint` is not configured.'
- )
-
- if not self.checkpoint_callback.best_model_path:
- if self.fast_dev_run:
- raise MisconfigurationException(
- f'You cannot execute `.{fn}(ckpt_path="best")` with `fast_dev_run=True`.'
- f" Please pass an exact checkpoint path to `.{fn}(ckpt_path=...)`"
- )
- raise MisconfigurationException(
- f'`.{fn}(ckpt_path="best")` is set but `ModelCheckpoint` is not configured to save the best model.'
- )
- # load best weights
- ckpt_path = self.checkpoint_callback.best_model_path
-
- if not ckpt_path:
- raise MisconfigurationException(
- f"`.{fn}()` found no path for the best weights: {ckpt_path!r}. Please"
- f" specify a path for a checkpoint `.{fn}(ckpt_path=PATH)`"
- )
- return ckpt_path
-
- def _call_setup_hook(self) -> None:
- fn = self.state.fn._setup_fn
-
- self.strategy.barrier("pre_setup")
-
- if self.datamodule is not None:
- self.datamodule.setup(stage=fn)
- self._call_callback_hooks("setup", stage=fn)
- self._call_lightning_module_hook("setup", stage=fn)
-
- self.strategy.barrier("post_setup")
-
- def _call_configure_sharded_model(self) -> None:
- with self.strategy.model_sharded_context():
- self._handle_meta_model()
- self._call_lightning_module_hook("configure_sharded_model")
- self._call_callback_hooks("on_configure_sharded_model")
-
- def _handle_meta_model(self) -> None:
- if not is_on_meta_device(self.lightning_module):
- return
-
- if isinstance(self.strategy, DDPSpawnStrategy):
- raise MisconfigurationException("LightningModule on meta device isn't supported with spawn.")
-
- materialize_module(self.lightning_module)
- # the trainer reference is lost during materialization
- self.lightning_module.trainer = proxy(self)
-
- def _call_teardown_hook(self) -> None:
- fn = self.state.fn._setup_fn
-
- if self.datamodule is not None:
- self.datamodule.teardown(stage=fn)
-
- self._call_callback_hooks("teardown", stage=fn)
- self._call_lightning_module_hook("teardown", stage=fn)
-
- self.lightning_module._current_fx_name = None
- # these could have become stale if metrics are defined in `setup`
- self.lightning_module._metric_attributes = None
-
- # todo: TPU 8 cores hangs in flush with TensorBoard. Might do for all loggers.
- # It might be related to xla tensors blocked when moving the cpu kill loggers.
- for logger in self.loggers:
- logger.finalize("success")
-
- # summarize profile results
- self.profiler.describe()
-
- def call_hook(
- self, hook_name: str, *args: Any, pl_module: Optional["pl.LightningModule"] = None, **kwargs: Any
- ) -> Any:
- r"""
- .. deprecated:: v1.6
- The Trainer's `call_hook` method was deprecated in v1.6 and will be removed in v1.8.
- """
- rank_zero_deprecation("The Trainer's `call_hook` method was deprecated in v1.6 and will be removed in v1.8.")
- pl_module = self.lightning_module or pl_module
- if pl_module:
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = hook_name
-
- # always profile hooks
- with self.profiler.profile(hook_name):
-
- # first call trainer hook
- callback_fx = getattr(self, hook_name, None)
- if callable(callback_fx):
- callback_fx(*args, **kwargs)
-
- # next call hook in lightningModule
- output = None
- model_fx = getattr(pl_module, hook_name, None)
- if callable(model_fx):
- output = model_fx(*args, **kwargs)
-
- # call the strategy hook
- if hook_name not in ("setup", "teardown", "on_train_start") and hasattr(self.strategy, hook_name):
- strategy_hook = getattr(self.strategy, hook_name)
- strategy_output = strategy_hook(*args, **kwargs)
- output = strategy_output if output is None else output
-
- if pl_module:
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- return output
-
- def _call_lightning_module_hook(
- self,
- hook_name: str,
- *args: Any,
- pl_module: Optional["pl.LightningModule"] = None,
- **kwargs: Any,
- ) -> Any:
- pl_module = pl_module or self.lightning_module
-
- if pl_module is None:
- raise TypeError("No Lightning Module is available to call hooks on")
-
- fn = getattr(pl_module, hook_name)
- if not callable(fn):
- return
-
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = hook_name
-
- with self.profiler.profile(f"[LightningModule]{pl_module.__class__.__name__}.{hook_name}"):
- output = fn(*args, **kwargs)
-
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- return output
-
- def _call_callback_hooks(
- self,
- hook_name: str,
- *args: Any,
- **kwargs: Any,
- ) -> None:
- log.debug(f"{self.__class__.__name__}: calling callback hook: {hook_name}")
- # TODO: remove if block in v1.8
- if hook_name in ("on_init_start", "on_init_end"):
- # these `Callback` hooks are the only ones that do not take a lightning module.
- # we also don't profile bc profiler hasn't been set yet
- for callback in self.callbacks:
- fn = getattr(callback, hook_name)
- if callable(fn):
- fn(self, *args, **kwargs)
- return
-
- pl_module = self.lightning_module
- if pl_module:
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = hook_name
-
- # TODO: remove if block in v1.7
- if hook_name == "on_train_batch_start":
- with self.profiler.profile(hook_name):
- self._on_train_batch_start(*args, **kwargs)
- elif hook_name == "on_train_batch_end":
- with self.profiler.profile(hook_name):
- self._on_train_batch_end(*args, **kwargs)
- else:
- for callback in self.callbacks:
- fn = getattr(callback, hook_name)
- if callable(fn):
- with self.profiler.profile(f"[Callback]{callback.state_key}.{hook_name}"):
- fn(self, self.lightning_module, *args, **kwargs)
-
- if pl_module:
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- # TODO: Delete this in v1.7 (deprecations: #9816 and #11148)
- def _on_train_batch_start(self, batch, batch_idx, dataloader_idx=0):
- r"""Called when the training batch begins. This function is needed because of two different deprecations affecting
- the original function in TrainerCallbackHookMixin: #9816 and #11148.
- """
- for callback in self.callbacks:
- if is_param_in_hook_signature(callback.on_train_batch_start, "dataloader_idx", explicit=True):
- callback.on_train_batch_start(self, self.lightning_module, batch, batch_idx, 0)
- else:
- callback.on_train_batch_start(self, self.lightning_module, batch, batch_idx)
-
- # TODO: Delete this in v1.7 (deprecations: #9816 and #11148)
- def _on_train_batch_end(self, outputs: STEP_OUTPUT, batch, batch_idx, dataloader_idx=0):
- r"""Called when the training batch ends. This function is needed because of two different deprecations affecting
- the original function in TrainerCallbackHookMixin: #9816 and #11148.
- """
- for callback in self.callbacks:
- if is_param_in_hook_signature(callback.on_train_batch_end, "dataloader_idx", explicit=True):
- callback.on_train_batch_end(self, self.lightning_module, outputs, batch, batch_idx, 0)
- else:
- callback.on_train_batch_end(self, self.lightning_module, outputs, batch, batch_idx)
-
- def _call_callbacks_state_dict(self) -> Dict[str, dict]:
- """Called when saving a model checkpoint, calls and returns every callback's `state_dict`, keyed by
- `Callback.state_key`."""
- callback_state_dicts = {}
- for callback in self.callbacks:
- state_dict = callback.state_dict()
- if state_dict:
- callback_state_dicts[callback.state_key] = state_dict
- return callback_state_dicts
-
- def _call_callbacks_on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
- """Called when saving a model checkpoint, calls every callback's `on_save_checkpoint` hook.
-
- Will be removed in v1.8: If state is returned, we insert the callback state into
- ``checkpoint["callbacks"][Callback.state_key]``. It overrides ``state_dict`` if already present.
- """
- pl_module = self.lightning_module
- if pl_module:
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = "on_save_checkpoint"
-
- for callback in self.callbacks:
- with self.profiler.profile(f"[Callback]{callback.state_key}.on_save_checkpoint"):
- state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
- if state:
- rank_zero_deprecation(
- f"Returning a value from `{callback.__class__.__name__}.on_save_checkpoint` is deprecated in v1.6"
- " and will be removed in v1.8. Please override `Callback.state_dict`"
- " to return state to be saved."
- )
- checkpoint["callbacks"][callback.state_key] = state
-
- if pl_module:
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- def _call_callbacks_on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
- """Called when loading a model checkpoint.
-
- Calls every callback's `on_load_checkpoint` hook. We have a dedicated function for this rather than using
- `_call_callback_hooks` because we have special logic for getting callback_states.
- """
- pl_module = self.lightning_module
- if pl_module:
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = "on_load_checkpoint"
-
- callback_states: Dict[Union[Type, str], Dict] = checkpoint.get("callbacks")
-
- if callback_states is None:
- return
-
- is_legacy_ckpt = Version(checkpoint["pytorch-lightning_version"]) < Version("1.5.0dev")
- current_callbacks_keys = {cb._legacy_state_key if is_legacy_ckpt else cb.state_key for cb in self.callbacks}
- difference = callback_states.keys() - current_callbacks_keys
- if difference:
- rank_zero_warn(
- "Be aware that when using `ckpt_path`,"
- " callbacks used to create the checkpoint need to be provided during `Trainer` instantiation."
- f" Please add the following callbacks: {list(difference)}.",
- )
-
- for callback in self.callbacks:
- state = callback_states.get(callback.state_key, callback_states.get(callback._legacy_state_key))
- if state:
- state = deepcopy(state)
- with self.profiler.profile(f"[Callback]{callback.state_key}.on_load_checkpoint"):
- callback.on_load_checkpoint(self, self.lightning_module, state)
-
- if pl_module:
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- def _call_callbacks_load_state_dict(self, checkpoint: Dict[str, Any]) -> None:
- """Called when loading a model checkpoint, calls every callback's `load_state_dict`."""
- callback_states: Dict[Union[Type, str], Dict] = checkpoint.get("callbacks")
-
- if callback_states is None:
- return
-
- for callback in self.callbacks:
- state = callback_states.get(callback.state_key, callback_states.get(callback._legacy_state_key))
- if state:
- state = deepcopy(state)
- callback.load_state_dict(state)
-
- def _call_strategy_hook(
- self,
- hook_name: str,
- *args: Any,
- **kwargs: Any,
- ) -> Any:
- pl_module = self.lightning_module
- prev_fx_name = pl_module._current_fx_name
- pl_module._current_fx_name = hook_name
-
- fn = getattr(self.strategy, hook_name)
- if not callable(fn):
- return
-
- with self.profiler.profile(f"[Strategy]{self.strategy.__class__.__name__}.{hook_name}"):
- output = fn(*args, **kwargs)
-
- # restore current_fx when nested context
- pl_module._current_fx_name = prev_fx_name
-
- return output
-
- @staticmethod
- def _log_api_event(event: str) -> None:
- torch._C._log_api_usage_once("lightning.trainer." + event)
-
- def __init_profiler(self, profiler: Optional[Union[Profiler, str]]) -> None:
- if isinstance(profiler, str):
- PROFILERS = {
- "simple": SimpleProfiler,
- "advanced": AdvancedProfiler,
- "pytorch": PyTorchProfiler,
- "xla": XLAProfiler,
- }
- profiler = profiler.lower()
- if profiler not in PROFILERS:
- raise MisconfigurationException(
- "When passing string value for the `profiler` parameter of `Trainer`,"
- f" it can only be one of {list(PROFILERS.keys())}"
- )
- profiler_class = PROFILERS[profiler]
- profiler = profiler_class()
- self.profiler: Profiler = profiler or PassThroughProfiler()
-
- def __setup_profiler(self) -> None:
- local_rank = self.local_rank if self.world_size > 1 else None
- self.profiler._lightning_module = proxy(self.lightning_module)
- self.profiler.setup(stage=self.state.fn._setup_fn, local_rank=local_rank, log_dir=self.log_dir)
-
- def _log_device_info(self) -> None:
- rank_zero_info(
- f"GPU available: {torch.cuda.is_available()}, used: {isinstance(self.accelerator, GPUAccelerator)}"
- )
-
- num_tpu_cores = self.num_devices if isinstance(self.accelerator, TPUAccelerator) else 0
- rank_zero_info(f"TPU available: {_TPU_AVAILABLE}, using: {num_tpu_cores} TPU cores")
-
- num_ipus = self.num_devices if isinstance(self.accelerator, IPUAccelerator) else 0
- rank_zero_info(f"IPU available: {_IPU_AVAILABLE}, using: {num_ipus} IPUs")
-
- num_hpus = self.num_devices if isinstance(self.accelerator, HPUAccelerator) else 0
- rank_zero_info(f"HPU available: {_HPU_AVAILABLE}, using: {num_hpus} HPUs")
-
- if torch.cuda.is_available() and not isinstance(self.accelerator, GPUAccelerator):
- rank_zero_warn(
- "GPU available but not used. Set `accelerator` and `devices` using"
- f" `Trainer(accelerator='gpu', devices={GPUAccelerator.auto_device_count()})`.",
- category=PossibleUserWarning,
- )
-
- if _TPU_AVAILABLE and not isinstance(self.accelerator, TPUAccelerator):
- rank_zero_warn(
- "TPU available but not used. Set `accelerator` and `devices` using"
- f" `Trainer(accelerator='tpu', devices={TPUAccelerator.auto_device_count()})`."
- )
-
- if _IPU_AVAILABLE and not isinstance(self.accelerator, IPUAccelerator):
- rank_zero_warn(
- "IPU available but not used. Set `accelerator` and `devices` using"
- f" `Trainer(accelerator='ipu', devices={IPUAccelerator.auto_device_count()})`."
- )
-
- if _HPU_AVAILABLE and not isinstance(self.accelerator, HPUAccelerator):
- rank_zero_warn(
- "HPU available but not used. Set `accelerator` and `devices` using"
- f" `Trainer(accelerator='hpu', devices={HPUAccelerator.auto_device_count()})`."
- )
-
- """
- Data loading methods
- """
-
- def reset_train_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
- """Resets the train dataloader and initialises required variables (number of batches, when to validate,
- etc.).
-
- Args:
- model: The ``LightningModule`` if calling this outside of the trainer scope.
- """
- source = self._data_connector._train_dataloader_source
- pl_module = self.lightning_module or model
- has_step = is_overridden("training_step", pl_module)
- enable_training = self.limit_train_batches > 0
- if not (source.is_defined() and has_step and enable_training):
- return
-
- self.train_dataloader = self._data_connector._request_dataloader(RunningStage.TRAINING, model=model)
-
- if self.overfit_batches > 0:
- self.train_dataloader = self._data_connector._resolve_overfit_batches(self.train_dataloader)
-
- # automatically add samplers
- self.train_dataloader = apply_to_collection(
- self.train_dataloader,
- (DataLoader, CombinedLoader),
- self._data_connector._prepare_dataloader,
- mode=RunningStage.TRAINING,
- )
- loaders = (
- self.train_dataloader.loaders
- if isinstance(self.train_dataloader, CombinedLoader)
- else self.train_dataloader
- )
-
- # check the workers recursively
- apply_to_collection(loaders, DataLoader, self._data_connector._worker_check, "train_dataloader")
-
- # add worker_init_fn for correct seeding in worker processes
- apply_to_collection(loaders, DataLoader, _auto_add_worker_init_fn, rank=self.global_rank)
-
- # add collate_fn to collect metadata for fault tolerant training
- if _fault_tolerant_training():
- apply_to_collection(loaders, DataLoader, _add_capture_metadata_collate)
-
- # wrap the sequence of train loaders to a CombinedLoader object for computing the num_training_batches
- if not isinstance(self.train_dataloader, CombinedLoader):
- self.train_dataloader = CombinedLoader(loaders, self._data_connector.multiple_trainloader_mode)
-
- module = model or self.lightning_module or self.datamodule
- self.num_training_batches = (
- len(self.train_dataloader)
- if has_len_all_ranks(self.train_dataloader, self.strategy, module)
- else float("inf")
- )
-
- if isinstance(self.limit_train_batches, int):
- self.num_training_batches = min(self.num_training_batches, int(self.limit_train_batches))
- elif self.num_training_batches != float("inf"):
- self.num_training_batches = int(self.num_training_batches * self.limit_train_batches)
- elif self.limit_train_batches != 1.0:
- raise MisconfigurationException(
- "When using an IterableDataset for `limit_train_batches`,"
- " `Trainer(limit_train_batches)` must be `1.0` or an int. An int k specifies"
- " `num_training_batches` to use."
- )
-
- if isinstance(self.val_check_interval, int):
- self.val_check_batch = self.val_check_interval
- if self.val_check_batch > self.num_training_batches:
- raise ValueError(
- f"`val_check_interval` ({self.val_check_interval}) must be less than or equal "
- f"to the number of the training batches ({self.num_training_batches}). "
- "If you want to disable validation set `limit_val_batches` to 0.0 instead."
- )
- else:
- if not has_len_all_ranks(self.train_dataloader, self.strategy, module):
- if self.val_check_interval == 1.0:
- self.val_check_batch = float("inf")
- else:
- raise MisconfigurationException(
- "When using an IterableDataset for `train_dataloader`,"
- " `Trainer(val_check_interval)` must be `1.0` or an int. An int k specifies"
- " checking validation every k training batches."
- )
- else:
- self.val_check_batch = int(self.num_training_batches * self.val_check_interval)
- self.val_check_batch = max(1, self.val_check_batch)
-
- if self.loggers and self.num_training_batches < self.log_every_n_steps:
- rank_zero_warn(
- f"The number of training batches ({self.num_training_batches}) is smaller than the logging interval"
- f" Trainer(log_every_n_steps={self.log_every_n_steps}). Set a lower value for log_every_n_steps if"
- " you want to see logs for the training epoch.",
- category=PossibleUserWarning,
- )
-
- # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
- self._last_train_dl_reload_epoch = self.current_epoch
-
- def reset_val_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
- """Resets the validation dataloader and determines the number of batches.
-
- Args:
- model: The ``LightningModule`` if called outside of the trainer scope.
- """
- source = self._data_connector._val_dataloader_source
- pl_module = self.lightning_module or model
- has_step = is_overridden("validation_step", pl_module)
- enable_validation = self.limit_val_batches > 0
- if source.is_defined() and has_step and enable_validation:
- self.num_val_batches, self.val_dataloaders = self._data_connector._reset_eval_dataloader(
- RunningStage.VALIDATING, model=pl_module
- )
-
- # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
- self._last_val_dl_reload_epoch = self.current_epoch
-
- def reset_test_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
- """Resets the test dataloader and determines the number of batches.
-
- Args:
- model: The ``LightningModule`` if called outside of the trainer scope.
- """
- source = self._data_connector._test_dataloader_source
- pl_module = self.lightning_module or model
- has_step = is_overridden("test_step", pl_module)
- enable_testing = self.limit_test_batches > 0
- if source.is_defined() and has_step and enable_testing:
- self.num_test_batches, self.test_dataloaders = self._data_connector._reset_eval_dataloader(
- RunningStage.TESTING, model=pl_module
- )
-
- def reset_predict_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
- """Resets the predict dataloader and determines the number of batches.
-
- Args:
- model: The ``LightningModule`` if called outside of the trainer scope.
- """
- source = self._data_connector._predict_dataloader_source
- pl_module = self.lightning_module or model
- enable_prediction = self.limit_predict_batches > 0
- if source.is_defined() and enable_prediction:
- self.num_predict_batches, self.predict_dataloaders = self._data_connector._reset_eval_dataloader(
- RunningStage.PREDICTING, model=pl_module
- )
-
- def reset_train_val_dataloaders(self, model: Optional["pl.LightningModule"] = None) -> None:
- """Resets train and val dataloaders if none are attached to the trainer.
-
- The val dataloader must be initialized before training loop starts, as the training loop
- inspects the val dataloader to determine whether to run the evaluation loop.
- Args:
- model: The ``LightningModule`` if called outside of the trainer scope.
- """
- if self.train_dataloader is None:
- self.reset_train_dataloader(model=model)
- if self.val_dataloaders is None:
- self.reset_val_dataloader(model=model)
-
- """
- Accelerator properties
- """
-
- @property
- def accelerator(self) -> Accelerator:
- return self.strategy.accelerator
-
- @property
- def strategy(self) -> Strategy:
- return self._accelerator_connector.strategy
-
- @property
- def training_type_plugin(self) -> Strategy:
- rank_zero_deprecation(
- "`Trainer.training_type_plugin` is deprecated in v1.6 and will be removed in v1.8. Use"
- " `Trainer.strategy` instead."
- )
- return self.strategy
-
- @property
- def precision_plugin(self) -> PrecisionPlugin:
- return self.strategy.precision_plugin
-
- @property
- def global_rank(self) -> int:
- return self.strategy.global_rank
-
- @property
- def local_rank(self) -> int:
- # some strategies define a local rank
- return getattr(self.strategy, "local_rank", 0)
-
- @property
- def node_rank(self) -> int:
- # some strategies define a node rank
- return getattr(self.strategy, "node_rank", 0)
-
- @property
- def world_size(self) -> int:
- # some strategies define a world size
- return getattr(self.strategy, "world_size", 1)
-
- @property
- def should_rank_save_checkpoint(self) -> bool:
- rank_zero_deprecation(
- "`Trainer.should_rank_save_checkpoint` is deprecated in v1.6 and will be removed in v1.8.", stacklevel=5
- )
- strategy = self.strategy
- return (
- isinstance(strategy, pl.strategies.TPUSpawnStrategy) and strategy.local_rank == 0 or strategy.is_global_zero
- )
-
- @property
- def num_nodes(self) -> int:
- return getattr(self.strategy, "num_nodes", 1)
-
- @property
- def device_ids(self) -> List[int]:
- """List of device indexes per node."""
- devices = (
- self.strategy.parallel_devices
- if isinstance(self.strategy, ParallelStrategy)
- else [self.strategy.root_device]
- )
- device_ids = []
- for idx, device in enumerate(devices):
- if isinstance(device, torch.device):
- device_ids.append(device.index or idx)
- elif isinstance(device, int):
- device_ids.append(device)
- return device_ids
-
- @property
- def num_devices(self) -> int:
- """Number of devices the trainer uses per node."""
- return len(self.device_ids)
-
- @property
- def num_processes(self) -> int:
- rank_zero_deprecation(
- "`Trainer.num_processes` is deprecated in v1.6 and will be removed in v1.8. "
- "Please use `Trainer.num_devices` instead."
- )
- return self.num_devices
-
- @property
- def root_gpu(self) -> Optional[int]:
- rank_zero_deprecation(
- "`Trainer.root_gpu` is deprecated in v1.6 and will be removed in v1.8. "
- "Please use `Trainer.strategy.root_device.index` instead."
- )
- return self.strategy.root_device.index if isinstance(self.accelerator, GPUAccelerator) else None
-
- @property
- def tpu_cores(self) -> int:
- rank_zero_deprecation(
- "`Trainer.tpu_cores` is deprecated in v1.6 and will be removed in v1.8. "
- "Please use `Trainer.num_devices` instead."
- )
- return self.num_devices if isinstance(self.accelerator, TPUAccelerator) else 0
-
- @property
- def ipus(self) -> int:
- rank_zero_deprecation(
- "`Trainer.ipus` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.num_devices` instead."
- )
- return self.num_devices if isinstance(self.accelerator, IPUAccelerator) else 0
-
- @property
- def num_gpus(self) -> int:
- rank_zero_deprecation(
- "`Trainer.num_gpus` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.num_devices` instead."
- )
- return self.num_devices if isinstance(self.accelerator, GPUAccelerator) else 0
-
- @property
- def devices(self) -> int:
- rank_zero_deprecation(
- "`Trainer.devices` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.num_devices` or `Trainer.device_ids` to get device information instead."
- )
- return self.num_devices
-
- @property
- def data_parallel_device_ids(self) -> Optional[List[int]]:
- rank_zero_deprecation(
- "`Trainer.data_parallel_device_ids` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.device_ids` instead."
- )
- return self.device_ids if isinstance(self.accelerator, GPUAccelerator) else None
-
- @property
- def lightning_module(self) -> "pl.LightningModule":
- # TODO: this is actually an optional return
- return self.strategy.lightning_module
-
- @property
- def optimizers(self) -> List[Optimizer]:
- return self.strategy.optimizers
-
- @optimizers.setter
- def optimizers(self, new_optims: Optional[List[Optimizer]]) -> None:
- self.strategy.optimizers = new_optims
-
- @property
- def lightning_optimizers(self) -> Dict[int, LightningOptimizer]:
- rank_zero_deprecation(
- "`Trainer.lightning_optimizers` is deprecated in v1.6 and will be removed in v1.8", stacklevel=5
- )
- return self.strategy._lightning_optimizers
-
- @property
- def lr_scheduler_configs(self) -> List[LRSchedulerConfig]:
- return self.strategy.lr_scheduler_configs
-
- @property
- def lr_schedulers(self) -> List[Dict[str, Any]]:
- rank_zero_deprecation(
- "`Trainer.lr_schedulers` is deprecated in v1.6 and will be removed in v1.8."
- " You can use `trainer.lr_scheduler_configs` instead which contains dataclasses instead of dictionaries.",
- stacklevel=5,
- )
- from dataclasses import asdict
-
- return [asdict(config) for config in self.strategy.lr_scheduler_configs]
-
- @property
- def optimizer_frequencies(self) -> List[int]:
- return self.strategy.optimizer_frequencies
-
- @optimizer_frequencies.setter
- def optimizer_frequencies(self, new_freqs: List[int]) -> None:
- self.strategy.optimizer_frequencies = new_freqs
-
- @property
- def amp_backend(self) -> Optional[AMPType]:
- if isinstance(self.precision_plugin, ApexMixedPrecisionPlugin):
- return AMPType.APEX
- if isinstance(self.precision_plugin, NativeMixedPrecisionPlugin):
- return AMPType.NATIVE
- return None
-
- @property
- def precision(self) -> Union[str, int]:
- return self.strategy.precision_plugin.precision
-
- @property
- def scaler(self) -> Optional[Any]:
- return getattr(self.precision_plugin, "scaler", None)
-
- @property
- def gpus(self) -> Optional[Union[List[int], str, int]]:
- rank_zero_deprecation(
- "`Trainer.gpus` was deprecated in v1.6 and will be removed in v1.8."
- " Please use `Trainer.num_devices` or `Trainer.device_ids` to get device information instead."
- )
- return self._accelerator_connector.gpus
-
- @property
- def model(self) -> torch.nn.Module:
- """The LightningModule, but possibly wrapped into DataParallel or DistributedDataParallel.
-
- To access the pure LightningModule, use
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.lightning_module` instead.
- """
- return self.strategy.model
-
- @model.setter
- def model(self, model: torch.nn.Module) -> None:
- """Setter for the model, pass-through to accelerator and plugin where the model reference is stored. Used
- by the Tuner to reset the state of Trainer and Accelerator.
-
- Args:
- model: The LightningModule, possibly wrapped into DataParallel or DistributedDataParallel, depending
- on the backend.
- """
- self.strategy.model = model
-
- """
- General properties
- """
-
- @property
- def log_dir(self) -> Optional[str]:
- if len(self.loggers) == 1:
- if isinstance(self.logger, TensorBoardLogger):
- dirpath = self.logger.log_dir
- else:
- dirpath = self.logger.save_dir
- else:
- dirpath = self.default_root_dir
-
- dirpath = self.strategy.broadcast(dirpath)
- return dirpath
-
- @property
- def use_amp(self) -> bool:
- rank_zero_deprecation(
- "`Trainer.use_amp` is deprecated in v1.6.0 and will be removed in v1.8.0."
- " Please use `Trainer.amp_backend` instead."
- )
- return self.precision == 16
-
- @property
- def is_global_zero(self) -> bool:
- return self.strategy.is_global_zero
-
- @property
- def slurm_job_id(self) -> Optional[int]:
- rank_zero_deprecation("Method `slurm_job_id` is deprecated in v1.6.0 and will be removed in v1.7.0.")
- return SLURMEnvironment.job_id()
-
- @property
- def distributed_sampler_kwargs(self) -> Optional[dict]:
- if isinstance(self.strategy, ParallelStrategy):
- return self.strategy.distributed_sampler_kwargs
-
- @property
- def data_parallel(self) -> bool:
- return isinstance(self.strategy, ParallelStrategy)
-
- @property
- def progress_bar_dict(self) -> dict:
- """Read-only for progress bar metrics."""
- rank_zero_deprecation(
- "`trainer.progress_bar_dict` is deprecated in v1.5 and will be removed in v1.7."
- " Use `ProgressBarBase.get_metrics` instead."
- )
- ref_model = self.lightning_module
- ref_model = cast(pl.LightningModule, ref_model)
- if self.progress_bar_callback:
- return self.progress_bar_callback.get_metrics(self, ref_model)
- return self.progress_bar_metrics
-
- @property
- def enable_validation(self) -> bool:
- """Check if we should run validation during training."""
- return (
- self._data_connector._val_dataloader_source.is_defined()
- and is_overridden("validation_step", self.lightning_module)
- and self.limit_val_batches > 0
- )
-
- @property
- def default_root_dir(self) -> str:
- """The default location to save artifacts of loggers, checkpoints etc.
-
- It is used as a fallback if logger or checkpoint callback do not define specific save paths.
- """
- if get_filesystem(self._default_root_dir).protocol == "file":
- return os.path.normpath(self._default_root_dir)
- return self._default_root_dir
-
- @property
- def weights_save_path(self) -> str:
- """
- The default root location to save weights (checkpoints), e.g., when the
- :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint` does not define a file path.
-
- .. deprecated:: v1.6
- `Trainer.weights_save_path` has been deprecated in v1.6 and will be removed in v1.8.
- """
- rank_zero_deprecation("`Trainer.weights_save_path` has been deprecated in v1.6 and will be removed in v1.8.")
- return self._weights_save_path_internal
-
- # TODO: Remove _weights_save_path_internal in v1.8
- @property
- def _weights_save_path_internal(self) -> str:
- """This is an internal implementation of weights_save_path which allows weights_save_path to be used
- internally by the framework without emitting a deprecation warning.
-
- To be removed in v1.8.
- """
- if get_filesystem(self._weights_save_path).protocol == "file":
- return os.path.normpath(self._weights_save_path)
- return self._weights_save_path
-
- @property
- def early_stopping_callback(self) -> Optional[EarlyStopping]:
- """The first :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` callback in the
- Trainer.callbacks list, or ``None`` if it doesn't exist."""
- callbacks = self.early_stopping_callbacks
- return callbacks[0] if len(callbacks) > 0 else None
-
- @property
- def early_stopping_callbacks(self) -> List[EarlyStopping]:
- """A list of all instances of :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` found in
- the Trainer.callbacks list."""
- return [c for c in self.callbacks if isinstance(c, EarlyStopping)]
-
- @property
- def prediction_writer_callbacks(self) -> List[BasePredictionWriter]:
- """A list of all instances of :class:`~pytorch_lightning.callbacks.prediction_writer.BasePredictionWriter`
- found in the Trainer.callbacks list."""
- return [cb for cb in self.callbacks if isinstance(cb, BasePredictionWriter)]
-
- @property
- def checkpoint_callback(self) -> Optional[ModelCheckpoint]:
- """The first :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint` callback in the
- Trainer.callbacks list, or ``None`` if it doesn't exist."""
- callbacks = self.checkpoint_callbacks
- return callbacks[0] if len(callbacks) > 0 else None
-
- @property
- def checkpoint_callbacks(self) -> List[ModelCheckpoint]:
- """A list of all instances of :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint` found
- in the Trainer.callbacks list."""
- return [c for c in self.callbacks if isinstance(c, ModelCheckpoint)]
-
- @property
- def progress_bar_callback(self) -> Optional[ProgressBarBase]:
- """An instance of :class:`~pytorch_lightning.callbacks.progress.base.ProgressBarBase` found in the
- Trainer.callbacks list, or ``None`` if one doesn't exist."""
- for c in self.callbacks:
- if isinstance(c, ProgressBarBase):
- return c
- return None
-
- @property
- def resume_from_checkpoint(self) -> Optional[Union[str, Path]]:
- resume_from_checkpoint = self._checkpoint_connector.resume_from_checkpoint_fit_path
- if resume_from_checkpoint is not None:
- rank_zero_deprecation(
- "`trainer.resume_from_checkpoint` is deprecated in v1.5 and will be removed in v2.0."
- " Specify the fit checkpoint path with `trainer.fit(ckpt_path=)` instead.",
- stacklevel=5,
- )
-
- return resume_from_checkpoint
-
- @property
- def ckpt_path(self) -> Optional[str]:
- """Set to the path/URL of a checkpoint loaded via :meth:`~pytorch_lightning.trainer.trainer.Trainer.fit`,
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.validate`,
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.test`, or
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.predict`. ``None`` otherwise."""
- return self._ckpt_path
-
- @property
- def validated_ckpt_path(self) -> Optional[str]:
- rank_zero_deprecation(
- "The `Trainer.validated_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via"
- " `Trainer.ckpt_path` instead.",
- stacklevel=5,
- )
- return self._validated_ckpt_path
-
- @validated_ckpt_path.setter
- def validated_ckpt_path(self, ckpt_path: Optional[str]) -> None:
- rank_zero_deprecation(
- "The `Trainer.validated_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via the"
- " read-only `Trainer.ckpt_path`.",
- stacklevel=5,
- )
- self._validated_ckpt_path = ckpt_path
-
- @property
- def tested_ckpt_path(self) -> Optional[str]:
- rank_zero_deprecation(
- "The `Trainer.tested_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via"
- " `Trainer.ckpt_path` instead.",
- stacklevel=5,
- )
- return self._tested_ckpt_path
-
- @tested_ckpt_path.setter
- def tested_ckpt_path(self, ckpt_path: Optional[str]) -> None:
- rank_zero_deprecation(
- "The `Trainer.tested_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via the"
- " read-only `Trainer.ckpt_path` instead.",
- stacklevel=5,
- )
- self._tested_ckpt_path = ckpt_path
-
- @property
- def predicted_ckpt_path(self) -> Optional[str]:
- rank_zero_deprecation(
- "The `Trainer.predicted_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via"
- " `Trainer.ckpt_path` instead.",
- stacklevel=5,
- )
- return self._predicted_ckpt_path
-
- @predicted_ckpt_path.setter
- def predicted_ckpt_path(self, ckpt_path: Optional[str]) -> None:
- rank_zero_deprecation(
- "The `Trainer.predicted_ckpt_path` attribute was deprecated in v1.6 and will be removed in v1.8. The"
- " path of a checkpoint loaded via `Trainer.{fit,validate,test,predict}` should be accessed via the"
- " read-only `Trainer.ckpt_path` instead.",
- stacklevel=5,
- )
- self._predicted_ckpt_path = ckpt_path
-
- def save_checkpoint(
- self, filepath: _PATH, weights_only: bool = False, storage_options: Optional[Any] = None
- ) -> None:
- r"""
- Runs routine to create a checkpoint.
-
- Args:
- filepath: Path where checkpoint is saved.
- weights_only: If ``True``, will only save the model weights.
- storage_options: parameter for how to save to storage, passed to ``CheckpointIO`` plugin
-
- """
- self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
-
- """
- Parsing properties
- """
-
- @classmethod
- def default_attributes(cls) -> dict:
- init_signature = inspect.signature(cls)
- return {k: v.default for k, v in init_signature.parameters.items()}
-
- @classmethod
- def get_deprecated_arg_names(cls) -> List:
- """Returns a list with deprecated Trainer arguments."""
- depr_arg_names = []
- for name, val in cls.__dict__.items():
- if name.startswith("DEPRECATED") and isinstance(val, (tuple, list)):
- depr_arg_names.extend(val)
- return depr_arg_names
-
- @classmethod
- def from_argparse_args(cls: Any, args: Union[Namespace, ArgumentParser], **kwargs) -> Any:
- return from_argparse_args(cls, args, **kwargs)
-
- @classmethod
- def parse_argparser(cls, arg_parser: Union[ArgumentParser, Namespace]) -> Namespace:
- return parse_argparser(cls, arg_parser)
-
- @classmethod
- def match_env_arguments(cls) -> Namespace:
- return parse_env_variables(cls)
-
- @classmethod
- def add_argparse_args(cls, parent_parser: ArgumentParser, **kwargs) -> ArgumentParser:
- return add_argparse_args(cls, parent_parser, **kwargs)
-
- """
- State properties
- """
-
- @property
- def interrupted(self) -> bool:
- return self.state.status == TrainerStatus.INTERRUPTED
-
- @property
- def training(self) -> bool:
- return self.state.stage == RunningStage.TRAINING
-
- @training.setter
- def training(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.TRAINING
- elif self.training:
- self.state.stage = None
-
- @property
- def testing(self) -> bool:
- return self.state.stage == RunningStage.TESTING
-
- @testing.setter
- def testing(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.TESTING
- elif self.testing:
- self.state.stage = None
-
- @property
- def predicting(self) -> bool:
- return self.state.stage == RunningStage.PREDICTING
-
- @predicting.setter
- def predicting(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.PREDICTING
- elif self.predicting:
- self.state.stage = None
-
- @property
- def tuning(self) -> bool:
- return self.state.stage == RunningStage.TUNING
-
- @tuning.setter
- def tuning(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.TUNING
- elif self.tuning:
- self.state.stage = None
-
- @property
- def validating(self) -> bool:
- return self.state.stage == RunningStage.VALIDATING
-
- @validating.setter
- def validating(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.VALIDATING
- elif self.validating:
- self.state.stage = None
-
- @property
- def evaluating(self) -> bool:
- return self.state.stage and self.state.stage.evaluating
-
- @property
- def sanity_checking(self) -> bool:
- return self.state.stage == RunningStage.SANITY_CHECKING
-
- @sanity_checking.setter
- def sanity_checking(self, val: bool) -> None:
- if val:
- self.state.stage = RunningStage.SANITY_CHECKING
- elif self.sanity_checking:
- self.state.stage = None
-
- """
- Loop properties
- """
-
- @property
- def global_step(self) -> int:
- """The number of optimizer steps taken (does not reset each epoch).
-
- This includes multiple optimizers and TBPTT steps (if enabled).
- """
- return self.fit_loop.epoch_loop.global_step
-
- @property
- def current_epoch(self) -> int:
- """The current epoch, updated after the epoch end hooks are run."""
- return self.fit_loop.epoch_progress.current.completed
-
- @property
- def max_epochs(self) -> int:
- return self.fit_loop.max_epochs
-
- @property
- def min_epochs(self) -> int:
- return self.fit_loop.min_epochs
-
- @property
- def max_steps(self) -> int:
- return self.fit_loop.max_steps
-
- @property
- def min_steps(self) -> Optional[int]:
- return self.fit_loop.min_steps
-
- @property
- def is_last_batch(self) -> bool:
- return self.fit_loop.epoch_loop.batch_progress.is_last_batch
-
- @property
- def fit_loop(self) -> FitLoop:
- return self._fit_loop
-
- @fit_loop.setter
- def fit_loop(self, loop: FitLoop):
- """Attach a custom fit loop to this Trainer.
-
- It will run with
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.fit`.
- """
- loop.trainer = self
- self._fit_loop = loop
-
- @property
- def validate_loop(self) -> EvaluationLoop:
- return self._validate_loop
-
- @validate_loop.setter
- def validate_loop(self, loop: EvaluationLoop):
- """Attach a custom validation loop to this Trainer.
-
- It will run with
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.validate`. Note that this loop is different from the one
- running during training inside the :meth:`pytorch_lightning.trainer.trainer.Trainer.fit` call.
- """
- loop.trainer = self
- self._validate_loop = loop
-
- @property
- def test_loop(self) -> EvaluationLoop:
- return self._test_loop
-
- @test_loop.setter
- def test_loop(self, loop: EvaluationLoop):
- """Attach a custom test loop to this Trainer.
-
- It will run with
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.test`.
- """
- loop.trainer = self
- self._test_loop = loop
-
- @property
- def predict_loop(self) -> PredictionLoop:
- return self._predict_loop
-
- @predict_loop.setter
- def predict_loop(self, loop: PredictionLoop):
- """Attach a custom prediction loop to this Trainer.
-
- It will run with
- :meth:`~pytorch_lightning.trainer.trainer.Trainer.predict`.
- """
- loop.trainer = self
- self._predict_loop = loop
-
- @property
- def verbose_evaluate(self) -> bool:
- rank_zero_deprecation(
- "The `Trainer.verbose_evaluate` property has been deprecated and will be removed in v1.8. The current value"
- " returned is the union of the validate and test loop values. You can choose which one to access with"
- " `trainer.{validate,test}_loop.verbose`.",
- stacklevel=5,
- )
- return self.validate_loop.verbose or self.test_loop.verbose
-
- @verbose_evaluate.setter
- def verbose_evaluate(self, verbose: bool) -> None:
- rank_zero_deprecation(
- "The `Trainer.verbose_evaluate` property has been deprecated and will be removed in v1.8. This will set"
- " the value for both trainer.{validate,test}_loop.verbose`.",
- stacklevel=5,
- )
- self.validate_loop.verbose = verbose
- self.test_loop.verbose = verbose
-
- @property
- def _evaluation_loop(self) -> EvaluationLoop:
- if self.state.fn in (TrainerFn.FITTING, TrainerFn.TUNING):
- return self.fit_loop.epoch_loop.val_loop
- if self.state.fn == TrainerFn.VALIDATING:
- return self.validate_loop
- if self.state.fn == TrainerFn.TESTING:
- return self.test_loop
- raise RuntimeError("The `Trainer._evaluation_loop` property isn't defined. Accessed outside of scope")
-
- @property
- def _active_loop(self) -> Optional[Union[FitLoop, EvaluationLoop, PredictionLoop]]:
- if self.training:
- return self.fit_loop
- if self.sanity_checking or self.evaluating:
- return self._evaluation_loop
- if self.predicting:
- return self.predict_loop
-
- """
- Logging properties
- """
-
- @property
- def logger(self) -> Optional[LightningLoggerBase]:
- if len(self.loggers) == 0:
- return None
- if len(self.loggers) == 1:
- return self.loggers[0]
- else:
- rank_zero_warn(
- "Using trainer.logger when Trainer is configured to use multiple loggers."
- " This behavior will change in v1.8 when LoggerCollection is removed, and"
- " trainer.logger will return the first logger in trainer.loggers"
- )
- with warnings.catch_warnings():
- warnings.simplefilter("ignore")
- return LoggerCollection(self.loggers)
-
- @logger.setter
- def logger(self, logger: Optional[LightningLoggerBase]) -> None:
- if not logger:
- self.loggers = []
- elif isinstance(logger, LoggerCollection):
- self.loggers = list(logger)
- else:
- self.loggers = [logger]
-
- @property
- def loggers(self) -> List[LightningLoggerBase]:
- return self._loggers
-
- @loggers.setter
- def loggers(self, loggers: Optional[List[LightningLoggerBase]]) -> None:
- self._loggers = loggers if loggers else []
-
- @property
- def callback_metrics(self) -> dict:
- return self._logger_connector.callback_metrics
-
- @property
- def logged_metrics(self) -> dict:
- return self._logger_connector.logged_metrics
-
- @property
- def progress_bar_metrics(self) -> dict:
- return self._logger_connector.progress_bar_metrics
-
- @property
- def _results(self) -> Optional[_ResultCollection]:
- active_loop = self._active_loop
- if active_loop is not None:
- return active_loop._results
-
- def _exit_gracefully_on_signal(self) -> None:
- if not _fault_tolerant_training() or not self._should_terminate_gracefully():
- return
- raise ExitGracefullyException(0)
-
- def _should_terminate_gracefully(self) -> bool:
- value = torch.tensor(int(self._terminate_gracefully), device=self.strategy.root_device)
- return self.strategy.reduce(value, reduce_op="sum") > 0
-
- @property
- def weights_summary(self) -> Optional[str]:
- rank_zero_deprecation("`Trainer.weights_summary` is deprecated in v1.5 and will be removed in v1.7.")
- return self._weights_summary
-
- @weights_summary.setter
- def weights_summary(self, val: Optional[str]) -> None:
- rank_zero_deprecation("Setting `Trainer.weights_summary` is deprecated in v1.5 and will be removed in v1.7.")
- self._weights_summary = val
-
- """
- Other
- """
-
- @property
- def estimated_stepping_batches(self) -> Union[int, float]:
- r"""
- Estimated stepping batches for the complete training inferred from DataLoaders, gradient
- accumulation factor and distributed setup.
-
- Examples::
-
- def configure_optimizers(self):
- optimizer = ...
- scheduler = torch.optim.lr_scheduler.OneCycleLR(
- optimizer, max_lr=1e-3, total_steps=self.trainer.estimated_stepping_batches
- )
- return [optimizer], [scheduler]
-
- """
- accumulation_scheduler = self.accumulation_scheduler
-
- if accumulation_scheduler.epochs != [0]:
- raise MisconfigurationException(
- "Estimated stepping batches cannot be computed with different"
- " `accumulate_grad_batches` at different epochs."
- )
-
- # infinite training
- if self.max_epochs == -1 and self.max_steps == -1:
- return float("inf")
-
- if self.train_dataloader is None:
- rank_zero_info("Loading `train_dataloader` to estimate number of stepping batches.")
- self.reset_train_dataloader()
-
- total_batches = self.num_training_batches
-
- # iterable dataset
- if total_batches == float("inf"):
- return self.max_steps
-
- self.accumulate_grad_batches = accumulation_scheduler.get_accumulate_grad_batches(self.current_epoch)
- effective_batch_size = self.accumulate_grad_batches
- max_estimated_steps = math.ceil(total_batches / effective_batch_size) * max(self.max_epochs, 1)
-
- max_estimated_steps = min(max_estimated_steps, self.max_steps) if self.max_steps != -1 else max_estimated_steps
- return max_estimated_steps
-
- @property
- def terminate_on_nan(self) -> bool:
- rank_zero_deprecation("`Trainer.terminate_on_nan` is deprecated in v1.5 and will be removed in 1.7.")
- return self._terminate_on_nan
-
- @terminate_on_nan.setter
- def terminate_on_nan(self, val: bool) -> None:
- rank_zero_deprecation(
- f"Setting `Trainer.terminate_on_nan = {val}` is deprecated in v1.5 and will be removed in 1.7."
- f" Please set `Trainer(detect_anomaly={val})` instead."
- )
- self._terminate_on_nan = val # : 212
-
-
-def _determine_batch_limits(batches: Optional[Union[int, float]], name: str) -> Union[int, float]:
- if batches is None:
- # batches is optional to know if the user passed a value so that we can show the above info messages only to the
- # users that set a value explicitly
- return 1.0
-
- # differentiating based on the type can be error-prone for users. show a message describing the chosen behaviour
- if isinstance(batches, int) and batches == 1:
- if name == "limit_train_batches":
- message = "1 batch per epoch will be used."
- elif name == "val_check_interval":
- message = "validation will run after every batch."
- else:
- message = "1 batch will be used."
- rank_zero_info(f"`Trainer({name}=1)` was configured so {message}")
- elif isinstance(batches, float) and batches == 1.0:
- if name == "limit_train_batches":
- message = "100% of the batches per epoch will be used."
- elif name == "val_check_interval":
- message = "validation will run at the end of the training epoch."
- else:
- message = "100% of the batches will be used."
- rank_zero_info(f"`Trainer({name}=1.0)` was configured so {message}.")
-
- if 0 <= batches <= 1:
- return batches
- if batches > 1 and batches % 1.0 == 0:
- return int(batches)
- raise MisconfigurationException(
- f"You have passed invalid value {batches} for {name}, it has to be in [0.0, 1.0] or an int."
- )
-To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.
-
-