pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14
"""LightningDataModule for loading DataLoaders with ease."""
15

16
import inspect
17
from typing import IO, Any, Dict, Iterable, Optional, Union, cast
18

19
from lightning_utilities import apply_to_collection
20
from torch.utils.data import DataLoader, Dataset, IterableDataset
21
from typing_extensions import Self
22

23
import lightning.pytorch as pl
24
from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
25
from lightning.pytorch.core.hooks import DataHooks
26
from lightning.pytorch.core.mixins import HyperparametersMixin
27
from lightning.pytorch.core.saving import _load_from_checkpoint
28
from lightning.pytorch.utilities.model_helpers import _restricted_classmethod
29
from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
30

31

32
class LightningDataModule(DataHooks, HyperparametersMixin):
33
    """A DataModule standardizes the training, val, test splits, data preparation and transforms. The main advantage is
34
    consistent data splits, data preparation and transforms across models.
35

36
    Example::
37

38
        import lightning as L
39
        import torch.utils.data as data
40
        from lightning.pytorch.demos.boring_classes import RandomDataset
41

42
        class MyDataModule(L.LightningDataModule):
43
            def prepare_data(self):
44
                # download, IO, etc. Useful with shared filesystems
45
                # only called on 1 GPU/TPU in distributed
46
                ...
47

48
            def setup(self, stage):
49
                # make assignments here (val/train/test split)
50
                # called on every process in DDP
51
                dataset = RandomDataset(1, 100)
52
                self.train, self.val, self.test = data.random_split(
53
                    dataset, [80, 10, 10], generator=torch.Generator().manual_seed(42)
54
                )
55

56
            def train_dataloader(self):
57
                return data.DataLoader(self.train)
58

59
            def val_dataloader(self):
60
                return data.DataLoader(self.val)
61

62
            def test_dataloader(self):
63
                return data.DataLoader(self.test)
64

65
            def teardown(self):
66
                # clean up state after the trainer stops, delete files...
67
                # called on every process in DDP
68
                ...
69

70
    """
71

72
    name: Optional[str] = None
73
    CHECKPOINT_HYPER_PARAMS_KEY = "datamodule_hyper_parameters"
74
    CHECKPOINT_HYPER_PARAMS_NAME = "datamodule_hparams_name"
75
    CHECKPOINT_HYPER_PARAMS_TYPE = "datamodule_hparams_type"
76

77
    def __init__(self) -> None:
78
        super().__init__()
79
        # Pointer to the trainer object
80
        self.trainer: Optional["pl.Trainer"] = None
81

82
    @classmethod
83
    def from_datasets(
84
        cls,
85
        train_dataset: Optional[Union[Dataset, Iterable[Dataset]]] = None,
86
        val_dataset: Optional[Union[Dataset, Iterable[Dataset]]] = None,
87
        test_dataset: Optional[Union[Dataset, Iterable[Dataset]]] = None,
88
        predict_dataset: Optional[Union[Dataset, Iterable[Dataset]]] = None,
89
        batch_size: int = 1,
90
        num_workers: int = 0,
91
        **datamodule_kwargs: Any,
92
    ) -> "LightningDataModule":
93
        r"""Create an instance from torch.utils.data.Dataset.
94

95
        Args:
96
            train_dataset: Optional dataset or iterable of datasets to be used for train_dataloader()
97
            val_dataset: Optional dataset or iterable of datasets to be used for val_dataloader()
98
            test_dataset: Optional dataset or iterable of datasets to be used for test_dataloader()
99
            predict_dataset: Optional dataset or iterable of datasets to be used for predict_dataloader()
100
            batch_size: Batch size to use for each dataloader. Default is 1. This parameter gets forwarded to the
101
                ``__init__`` if the datamodule has such a name defined in its signature.
102
            num_workers: Number of subprocesses to use for data loading. 0 means that the
103
                data will be loaded in the main process. Number of CPUs available. This parameter gets forwarded to the
104
                ``__init__`` if the datamodule has such a name defined in its signature.
105
            **datamodule_kwargs: Additional parameters that get passed down to the datamodule's ``__init__``.
106

107
        """
108

109
        def dataloader(ds: Dataset, shuffle: bool = False) -> DataLoader:
110
            shuffle &= not isinstance(ds, IterableDataset)
111
            return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True)
112

113
        def train_dataloader() -> TRAIN_DATALOADERS:
114
            return apply_to_collection(train_dataset, Dataset, dataloader, shuffle=True)
115

116
        def val_dataloader() -> EVAL_DATALOADERS:
117
            return apply_to_collection(val_dataset, Dataset, dataloader)
118

119
        def test_dataloader() -> EVAL_DATALOADERS:
120
            return apply_to_collection(test_dataset, Dataset, dataloader)
121

122
        def predict_dataloader() -> EVAL_DATALOADERS:
123
            return apply_to_collection(predict_dataset, Dataset, dataloader)
124

125
        candidate_kwargs = {"batch_size": batch_size, "num_workers": num_workers}
126
        accepted_params = inspect.signature(cls.__init__).parameters
127
        accepts_kwargs = any(param.kind == param.VAR_KEYWORD for param in accepted_params.values())
128
        if accepts_kwargs:
129
            special_kwargs = candidate_kwargs
130
        else:
131
            accepted_param_names = set(accepted_params)
132
            accepted_param_names.discard("self")
133
            special_kwargs = {k: v for k, v in candidate_kwargs.items() if k in accepted_param_names}
134

135
        datamodule = cls(**datamodule_kwargs, **special_kwargs)
136
        if train_dataset is not None:
137
            datamodule.train_dataloader = train_dataloader  # type: ignore[method-assign]
138
        if val_dataset is not None:
139
            datamodule.val_dataloader = val_dataloader  # type: ignore[method-assign]
140
        if test_dataset is not None:
141
            datamodule.test_dataloader = test_dataloader  # type: ignore[method-assign]
142
        if predict_dataset is not None:
143
            datamodule.predict_dataloader = predict_dataloader  # type: ignore[method-assign]
144
        return datamodule
145

146
    def state_dict(self) -> Dict[str, Any]:
147
        """Called when saving a checkpoint, implement to generate and save datamodule state.
148

149
        Returns:
150
            A dictionary containing datamodule state.
151

152
        """
153
        return {}
154

155
    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
156
        """Called when loading a checkpoint, implement to reload datamodule state given datamodule state_dict.
157

158
        Args:
159
            state_dict: the datamodule state returned by ``state_dict``.
160

161
        """
162
        pass
163

164
    @_restricted_classmethod
165
    def load_from_checkpoint(
166
        cls,
167
        checkpoint_path: Union[_PATH, IO],
168
        map_location: _MAP_LOCATION_TYPE = None,
169
        hparams_file: Optional[_PATH] = None,
170
        **kwargs: Any,
171
    ) -> Self:
172
        r"""Primary way of loading a datamodule from a checkpoint. When Lightning saves a checkpoint it stores the
173
        arguments passed to ``__init__``  in the checkpoint under ``"datamodule_hyper_parameters"``.
174

175
        Any arguments specified through \*\*kwargs will override args stored in ``"datamodule_hyper_parameters"``.
176

177
        Args:
178
            checkpoint_path: Path to checkpoint. This can also be a URL, or file-like object
179
            map_location:
180
                If your checkpoint saved a GPU model and you now load on CPUs
181
                or a different number of GPUs, use this to map to the new setup.
182
                The behaviour is the same as in :func:`torch.load`.
183
            hparams_file: Optional path to a ``.yaml`` or ``.csv`` file with hierarchical structure
184
                as in this example::
185

186
                    dataloader:
187
                        batch_size: 32
188

189
                You most likely won't need this since Lightning will always save the hyperparameters
190
                to the checkpoint.
191
                However, if your checkpoint weights don't have the hyperparameters saved,
192
                use this method to pass in a ``.yaml`` file with the hparams you'd like to use.
193
                These will be converted into a :class:`~dict` and passed into your
194
                :class:`LightningDataModule` for use.
195

196
                If your datamodule's ``hparams`` argument is :class:`~argparse.Namespace`
197
                and ``.yaml`` file has hierarchical structure, you need to refactor your datamodule to treat
198
                ``hparams`` as :class:`~dict`.
199
            \**kwargs: Any extra keyword args needed to init the datamodule. Can also be used to override saved
200
                hyperparameter values.
201

202
        Return:
203
            :class:`LightningDataModule` instance with loaded weights and hyperparameters (if available).
204

205
        Note:
206
            ``load_from_checkpoint`` is a **class** method. You must use your :class:`LightningDataModule`
207
            **class** to call it instead of the :class:`LightningDataModule` instance, or a
208
            ``TypeError`` will be raised.
209

210
        Example::
211

212
            # load weights without mapping ...
213
            datamodule = MyLightningDataModule.load_from_checkpoint('path/to/checkpoint.ckpt')
214

215
            # or load weights and hyperparameters from separate files.
216
            datamodule = MyLightningDataModule.load_from_checkpoint(
217
                'path/to/checkpoint.ckpt',
218
                hparams_file='/path/to/hparams_file.yaml'
219
            )
220

221
            # override some of the params with new values
222
            datamodule = MyLightningDataModule.load_from_checkpoint(
223
                PATH,
224
                batch_size=32,
225
                num_workers=10,
226
            )
227

228
        """
229
        loaded = _load_from_checkpoint(
230
            cls,  # type: ignore[arg-type]
231
            checkpoint_path,
232
            map_location=map_location,
233
            hparams_file=hparams_file,
234
            strict=None,
235
            **kwargs,
236
        )
237
        return cast(Self, loaded)
238