Skills › Scientific Agent Skills › pytorch-lightning › template_datamodule.py
template_datamodule.py
170 lines · skills/pytorch-lightning/scripts/template_datamodule.py at 5f2bc03
1"""Small, deterministic DataModule example; replace synthetic data with real splits.
3The 784-feature samples match TemplateLightningModule. Random labels are only a
4loop smoke test, never evidence of scientific accuracy. No files are downloaded.
5"""
7import lightning as L
8import torch
9from torch.utils.data import DataLoader, Dataset, Subset
12class CustomDataset(Dataset):
13 """Synthetic labelled vectors with an explicit seed and optional transform."""
15 def __init__(self, data_path, transform=None, *, seed=42, num_samples=1000):
16 self.data_path = data_path # Reserved for a real dataset implementation.
17 self.transform = transform
18 generator = torch.Generator().manual_seed(seed)
19 self.data = torch.randn(num_samples, 784, generator=generator)
20 self.labels = torch.randint(0, 10, (num_samples,), generator=generator)
22 def __len__(self):
23 return len(self.data)
25 def __getitem__(self, idx):
26 sample = self.data[idx]
27 if self.transform is not None:
28 sample = self.transform(sample)
29 return sample, self.labels[idx]
32class TemplateDataModule(L.LightningDataModule):
33 """Seeded train/validation split with separate evaluation transforms.
35 For real data, preserve sample IDs and preprocessing provenance in the
36 checkpoint; use group/time splits where observations are correlated.
37 """
39 def __init__(
40 self,
41 data_dir: str = "./data",
42 batch_size: int = 32,
43 num_workers: int = 0,
44 train_val_split: float = 0.8,
45 pin_memory: bool = False,
46 seed: int = 42,
47 num_samples: int = 1000,
48 ):
49 super().__init__()
50 if not 0 < train_val_split < 1:
51 raise ValueError("train_val_split must be between 0 and 1")
52 if num_samples < 2 or not 0 < int(train_val_split * num_samples) < num_samples:
53 raise ValueError("Both train and validation splits must be nonempty")
54 if batch_size < 1 or num_workers < 0:
55 raise ValueError("batch_size must be positive and num_workers nonnegative")
56 self.save_hyperparameters()
57 self.train_dataset = self.val_dataset = None
58 self.test_dataset = self.predict_dataset = None
60 def prepare_data(self):
61 """Download/cache here, if needed; do not assign distributed state.
63 Lightning calls this on local rank zero on each node by default. Set
64 prepare_data_per_node=False for one global writer with shared storage.
65 This synthetic example has no preparation or network access.
66 """
68 def _dataset(self, transform, seed_offset=0):
69 return CustomDataset(
70 self.hparams.data_dir,
71 transform=transform,
72 seed=self.hparams.seed + seed_offset,
73 num_samples=self.hparams.num_samples,
74 )
76 def setup(self, stage=None):
77 """Build per-process state, including standalone validate()."""
78 if stage not in (None, "fit", "validate", "test", "predict"):
79 raise ValueError(f"Unsupported stage: {stage}")
80 if stage in (None, "fit", "validate"):
81 indices = torch.randperm(
82 self.hparams.num_samples,
83 generator=torch.Generator().manual_seed(self.hparams.seed),
84 ).tolist()
85 boundary = int(self.hparams.train_val_split * len(indices))
86 # Two dataset views have identical underlying samples but different
87 # transforms; mutating one Subset's dataset must not augment validation.
88 self.train_dataset = Subset(
89 self._dataset(self._get_train_transforms()), indices[:boundary]
90 )
91 self.val_dataset = Subset(
92 self._dataset(self._get_test_transforms()), indices[boundary:]
93 )
94 if stage in (None, "test"):
95 self.test_dataset = self._dataset(self._get_test_transforms(), seed_offset=1)
96 if stage in (None, "predict"):
97 self.predict_dataset = self._dataset(self._get_test_transforms(), seed_offset=2)
99 def _get_train_transforms(self):
100 """Override with training-only augmentation, fit on training data only."""
101 return None
103 def _get_test_transforms(self):
104 """Override with deterministic transforms using training-fitted parameters."""
105 return None
107 def _loader(self, dataset, *, training=False):
108 if dataset is None:
109 raise RuntimeError("Call setup() for this stage before requesting a DataLoader")
110 return DataLoader(
111 dataset,
112 batch_size=self.hparams.batch_size,
113 shuffle=training,
114 num_workers=self.hparams.num_workers,
115 pin_memory=self.hparams.pin_memory,
116 persistent_workers=self.hparams.num_workers > 0,
117 drop_last=False,
118 )
120 def train_dataloader(self):
121 return self._loader(self.train_dataset, training=True)
123 def val_dataloader(self):
124 return self._loader(self.val_dataset)
126 def test_dataloader(self):
127 return self._loader(self.test_dataset)
129 def predict_dataloader(self):
130 return self._loader(self.predict_dataset)
132 def state_dict(self):
133 """Enough provenance to regenerate this synthetic dataset and its split."""
134 return {
135 key: self.hparams[key]
136 for key in ("train_val_split", "seed", "num_samples")
137 }
139 def load_state_dict(self, state_dict):
140 """Lightning restores DataModule state after setup; rebuild existing views."""
141 stages = []
142 if self.train_dataset is not None or self.val_dataset is not None:
143 stages.append("fit")
144 if self.test_dataset is not None:
145 stages.append("test")
146 if self.predict_dataset is not None:
147 stages.append("predict")
148 for key in ("train_val_split", "seed", "num_samples"):
149 if key in state_dict: # Accept older checkpoints containing just ratio.
150 self.hparams[key] = state_dict[key]
151 if not 0 < int(self.hparams.train_val_split * self.hparams.num_samples) < self.hparams.num_samples:
152 raise ValueError("Restored split must have nonempty train and validation sets")
153 for stage in stages:
154 self.setup(stage)
156 def teardown(self, stage=None):
157 if stage in (None, "fit", "validate"):
158 self.train_dataset = self.val_dataset = None
159 if stage in (None, "test"):
160 self.test_dataset = None
161 if stage in (None, "predict"):
162 self.predict_dataset = None
165if __name__ == "__main__":
166 dm = TemplateDataModule(batch_size=64)
167 dm.prepare_data()
168 dm.setup("fit")
169 print(f"[OK] Train/validation samples: {len(dm.train_dataset)}/{len(dm.val_dataset)}")
170 print(f"[OK] Batch shape: {tuple(next(iter(dm.train_dataloader()))[0].shape)}")