Skip to content

datamodule

MEDSDataModule

Bases: LightningDataModule, Module

LightningDataModule for the MEDS pytorch dataset.

TODO: Add documentation

A LightningDataModule implements 7 key methods:

```python def prepare_data(self): # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP). # Pre-process and loading data, save to disk, etc…

def setup(self, stage): # Things to do on every process in DDP. # Load data, set variables, etc…

def train_dataloader(self): # return train dataloader

def val_dataloader(self): # return validation dataloader

def test_dataloader(self): # return test dataloader

def predict_dataloader(self): # return predict dataloader

def teardown(self, stage): # Called on every process in DDP. # Clean up after fit or test. ```

This allows you to share a full dataset without explaining how to download, split, transform and process the data.

Read the docs: https://lightning.ai/docs/pytorch/latest/data/datamodule.html

Source code in meds_torch/data/datamodule.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
class MEDSDataModule(LightningDataModule, Module):
    """`LightningDataModule` for the MEDS pytorch dataset.

    TODO: Add documentation

    A `LightningDataModule` implements 7 key methods:

    ```python     def prepare_data(self):     # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP).     #
    Pre-process and loading data, save to disk, etc...

    def setup(self, stage): # Things to do on every process in DDP. # Load data, set variables, etc...

    def train_dataloader(self): # return train dataloader

    def val_dataloader(self): # return validation dataloader

    def test_dataloader(self): # return test dataloader

    def predict_dataloader(self): # return predict dataloader

    def teardown(self, stage):     # Called on every process in DDP.     # Clean up after fit or test. ```

    This allows you to share a full dataset without explaining how to download, split, transform and process
    the data.

    Read the docs:
    https://lightning.ai/docs/pytorch/latest/data/datamodule.html
    """

    def __init__(
        self,
        cfg: DictConfig = None,
    ) -> None:
        """Initialize a `MEDSDataModule`."""
        super().__init__()
        self.cfg = cfg

        # this line allows to access init params with 'self.hparams' attribute
        # also ensures init params will be stored in ckpt
        self.save_hyperparameters(logger=False)

        self.data_train: Dataset | None = None
        self.data_val: Dataset | None = None
        self.data_test: Dataset | None = None

    @property
    def num_classes(self) -> int:
        """Get the number of classes.

        :return: The number of MNIST classes (10).
        """
        return 10

    def prepare_data(self) -> None:
        """Download data if needed. Lightning ensures that `self.prepare_data()` is
        called only within a single process on CPU, so you can safely add your
        downloading logic within. In case of multi-node training, the execution of this
        hook depends upon `self.prepare_data_per_node()`.

        Do not use it to assign state (self.x = y).
        """

    def setup(self, stage: str | None = None) -> None:
        """Load data. Set variables: `self.data_train`, `self.data_val`,
        `self.data_test`.

        This method is called by Lightning before `trainer.fit()`, `trainer.validate()`, `trainer.test()`, and
        `trainer.predict()`, so be careful not to execute things like random split twice! Also, it is called
        after `self.prepare_data()` and there is a barrier in between which ensures that all the processes
        proceed to `self.setup()` once the data is prepared and available for use.

        :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to
            ``None``.
        """
        # Divide batch size by the number of devices.
        if self.trainer is not None:
            if self.hparams.cfg.dataloader.batch_size % self.trainer.world_size != 0:
                raise RuntimeError(
                    f"Batch size ({self.hparams.cfg.dataloader.batch_size}) is not divisible by "
                    f"the number of devices ({self.trainer.world_size})."
                )
            self.batch_size_per_device = self.hparams.cfg.dataloader.batch_size // self.trainer.world_size

        # load and split datasets only if not loaded already
        if stage == "test":
            self.data_test = get_dataset(self.cfg, split=self.cfg.split_names.test)
        elif stage == "validate":
            self.data_val = get_dataset(self.cfg, split=self.cfg.split_names.validate)
        else:
            self.data_train = get_dataset(self.cfg, split=self.cfg.split_names.train)
            self.data_val = get_dataset(self.cfg, split=self.cfg.split_names.validate)
            self.data_test = get_dataset(self.cfg, split=self.cfg.split_names.test)

    def train_dataloader(self) -> DataLoader[Any]:
        """Create and return the train dataloader.

        :return: The train dataloader.
        """
        return DataLoader(
            dataset=self.data_train,
            shuffle=True,
            collate_fn=self.data_train.collate,
            drop_last=True,
            **self.cfg.dataloader,
        )

    def val_dataloader(self) -> DataLoader[Any]:
        """Create and return the validation dataloader.

        :return: The validation dataloader.
        """
        return DataLoader(
            dataset=self.data_val,
            shuffle=False,
            collate_fn=self.data_val.collate,
            **self.cfg.dataloader,
        )

    def test_dataloader(self) -> DataLoader[Any]:
        """Create and return the test dataloader.

        :return: The test dataloader.
        """
        return DataLoader(
            dataset=self.data_test,
            shuffle=False,
            collate_fn=self.data_test.collate,
            **self.cfg.dataloader,
        )

    def predict_dataloader(self) -> DataLoader[Any]:
        """Create and return the predict dataloader.

        :return: The predict dataloader.
        """
        if self.cfg.predict_dataset == "train":
            return self.train_dataloader()
        elif self.cfg.predict_dataset == "val":
            return self.val_dataloader()
        elif self.cfg.predict_dataset == "test":
            return self.test_dataloader()
        else:
            raise NotImplementedError(
                f"{self.cfg.predict_dataset} not implemented! Use 'train', 'val', or 'test'."
            )

    def state_dict(self) -> dict[Any, Any]:
        """Called when saving a checkpoint. Implement to generate and save the
        datamodule state.

        :return: A dictionary containing the datamodule state that you want to save.
        """
        return {}

    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()`.

        :param state_dict: The datamodule state returned by `self.state_dict()`.
        """

num_classes: int property

Get the number of classes.

:return: The number of MNIST classes (10).

__init__(cfg=None)

Initialize a MEDSDataModule.

Source code in meds_torch/data/datamodule.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    cfg: DictConfig = None,
) -> None:
    """Initialize a `MEDSDataModule`."""
    super().__init__()
    self.cfg = cfg

    # this line allows to access init params with 'self.hparams' attribute
    # also ensures init params will be stored in ckpt
    self.save_hyperparameters(logger=False)

    self.data_train: Dataset | None = None
    self.data_val: Dataset | None = None
    self.data_test: Dataset | None = None

load_state_dict(state_dict)

Called when loading a checkpoint. Implement to reload datamodule state given datamodule state_dict().

:param state_dict: The datamodule state returned by self.state_dict().

Source code in meds_torch/data/datamodule.py 
170
171
172
173
174
175
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()`.

    :param state_dict: The datamodule state returned by `self.state_dict()`.
    """

predict_dataloader()

Create and return the predict dataloader.

:return: The predict dataloader.

Source code in meds_torch/data/datamodule.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def predict_dataloader(self) -> DataLoader[Any]:
    """Create and return the predict dataloader.

    :return: The predict dataloader.
    """
    if self.cfg.predict_dataset == "train":
        return self.train_dataloader()
    elif self.cfg.predict_dataset == "val":
        return self.val_dataloader()
    elif self.cfg.predict_dataset == "test":
        return self.test_dataloader()
    else:
        raise NotImplementedError(
            f"{self.cfg.predict_dataset} not implemented! Use 'train', 'val', or 'test'."
        )

prepare_data()

Download data if needed. Lightning ensures that self.prepare_data() is called only within a single process on CPU, so you can safely add your downloading logic within. In case of multi-node training, the execution of this hook depends upon self.prepare_data_per_node().

Do not use it to assign state (self.x = y).

Source code in meds_torch/data/datamodule.py 
69
70
71
72
73
74
75
76
def prepare_data(self) -> None:
    """Download data if needed. Lightning ensures that `self.prepare_data()` is
    called only within a single process on CPU, so you can safely add your
    downloading logic within. In case of multi-node training, the execution of this
    hook depends upon `self.prepare_data_per_node()`.

    Do not use it to assign state (self.x = y).
    """

setup(stage=None)

Load data. Set variables: self.data_train, self.data_val, self.data_test.

This method is called by Lightning before trainer.fit(), trainer.validate(), trainer.test(), and trainer.predict(), so be careful not to execute things like random split twice! Also, it is called after self.prepare_data() and there is a barrier in between which ensures that all the processes proceed to self.setup() once the data is prepared and available for use.

:param stage: The stage to setup. Either "fit", "validate", "test", or "predict". Defaults to None.

Source code in meds_torch/data/datamodule.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def setup(self, stage: str | None = None) -> None:
    """Load data. Set variables: `self.data_train`, `self.data_val`,
    `self.data_test`.

    This method is called by Lightning before `trainer.fit()`, `trainer.validate()`, `trainer.test()`, and
    `trainer.predict()`, so be careful not to execute things like random split twice! Also, it is called
    after `self.prepare_data()` and there is a barrier in between which ensures that all the processes
    proceed to `self.setup()` once the data is prepared and available for use.

    :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to
        ``None``.
    """
    # Divide batch size by the number of devices.
    if self.trainer is not None:
        if self.hparams.cfg.dataloader.batch_size % self.trainer.world_size != 0:
            raise RuntimeError(
                f"Batch size ({self.hparams.cfg.dataloader.batch_size}) is not divisible by "
                f"the number of devices ({self.trainer.world_size})."
            )
        self.batch_size_per_device = self.hparams.cfg.dataloader.batch_size // self.trainer.world_size

    # load and split datasets only if not loaded already
    if stage == "test":
        self.data_test = get_dataset(self.cfg, split=self.cfg.split_names.test)
    elif stage == "validate":
        self.data_val = get_dataset(self.cfg, split=self.cfg.split_names.validate)
    else:
        self.data_train = get_dataset(self.cfg, split=self.cfg.split_names.train)
        self.data_val = get_dataset(self.cfg, split=self.cfg.split_names.validate)
        self.data_test = get_dataset(self.cfg, split=self.cfg.split_names.test)

state_dict()

Called when saving a checkpoint. Implement to generate and save the datamodule state.

:return: A dictionary containing the datamodule state that you want to save.

Source code in meds_torch/data/datamodule.py 
162
163
164
165
166
167
168
def state_dict(self) -> dict[Any, Any]:
    """Called when saving a checkpoint. Implement to generate and save the
    datamodule state.

    :return: A dictionary containing the datamodule state that you want to save.
    """
    return {}

test_dataloader()

Create and return the test dataloader.

:return: The test dataloader.

Source code in meds_torch/data/datamodule.py 
134
135
136
137
138
139
140
141
142
143
144
def test_dataloader(self) -> DataLoader[Any]:
    """Create and return the test dataloader.

    :return: The test dataloader.
    """
    return DataLoader(
        dataset=self.data_test,
        shuffle=False,
        collate_fn=self.data_test.collate,
        **self.cfg.dataloader,
    )

train_dataloader()

Create and return the train dataloader.

:return: The train dataloader.

Source code in meds_torch/data/datamodule.py 
109
110
111
112
113
114
115
116
117
118
119
120
def train_dataloader(self) -> DataLoader[Any]:
    """Create and return the train dataloader.

    :return: The train dataloader.
    """
    return DataLoader(
        dataset=self.data_train,
        shuffle=True,
        collate_fn=self.data_train.collate,
        drop_last=True,
        **self.cfg.dataloader,
    )

val_dataloader()

Create and return the validation dataloader.

:return: The validation dataloader.

Source code in meds_torch/data/datamodule.py 
122
123
124
125
126
127
128
129
130
131
132
def val_dataloader(self) -> DataLoader[Any]:
    """Create and return the validation dataloader.

    :return: The validation dataloader.
    """
    return DataLoader(
        dataset=self.data_val,
        shuffle=False,
        collate_fn=self.data_val.collate,
        **self.cfg.dataloader,
    )