MNISTDataModule
MNISTDataModule (data_dir:str='~/Data/',
train_val_test_split:List[float]=[0.8, 0.1, 0.1],
batch_size:int=64, num_workers:int=0,
pin_memory:bool=False, persistent_workers:bool=False)
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::
import lightning.pytorch as L
import torch.utils.data as data
from pytorch_lightning.demos.boring_classes import RandomDataset
class MyDataModule(L.LightningDataModule):
def prepare_data(self):
# download, IO, etc. Useful with shared filesystems
# 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
dataset = RandomDataset(1, 100)
self.train, self.val, self.test = data.random_split(
dataset, [80, 10, 10], generator=torch.Generator().manual_seed(42)
)
def train_dataloader(self):
return data.DataLoader(self.train)
def val_dataloader(self):
return data.DataLoader(self.val)
def test_dataloader(self):
return data.DataLoader(self.test)
def teardown(self):
# clean up state after the trainer stops, delete files...
# called on every process in DDP
...
| data_dir |
str |
~/Data/ |
path to source data dir |
| train_val_test_split |
typing.List[float] |
[0.8, 0.1, 0.1] |
train val test % |
| batch_size |
int |
64 |
size of compute batch |
| num_workers |
int |
0 |
num_workers equal 0 means that it’s the main process that will do the data loading when needed, num_workers equal 1 is the same as any n, but you’ll only have a single worker, so it might be slow |
| pin_memory |
bool |
False |
If you load your samples in the Dataset on CPU and would like to push it during training to the GPU, you can speed up the host to device transfer by enabling pin_memory. This lets your DataLoader allocate the samples in page-locked memory, which speeds-up the transfer |
| persistent_workers |
bool |
False |
|
