Skills › Scientific Agent Skills › pytorch-lightning › template_lightning_module.py
template_lightning_module.py
220 lines · skills/pytorch-lightning/scripts/template_lightning_module.py at 5f2bc03
1"""
2Template for creating a PyTorch Lightning Module.
4This template provides a complete boilerplate for building a LightningModule
5with all essential methods and best practices.
6"""
8import lightning as L
9import torch
10import torch.nn as nn
11import torch.nn.functional as F
12from torch.optim import Adam
13from torch.optim.lr_scheduler import ReduceLROnPlateau
16class TemplateLightningModule(L.LightningModule):
17 """
18 Template LightningModule for building deep learning models.
20 Args:
21 learning_rate: Learning rate for optimizer
22 hidden_dim: Hidden dimension size
23 dropout: Dropout probability
24 """
26 def __init__(
27 self,
28 learning_rate: float = 0.001,
29 hidden_dim: int = 256,
30 dropout: float = 0.1,
31 ):
32 super().__init__()
34 # Save hyperparameters (accessible via self.hparams)
35 self.save_hyperparameters()
37 # Define your model architecture
38 self.model = nn.Sequential(
39 nn.Linear(784, self.hparams.hidden_dim),
40 nn.ReLU(),
41 nn.Dropout(self.hparams.dropout),
42 nn.Linear(self.hparams.hidden_dim, 10),
43 )
45 # Optional: Define metrics
46 # from torchmetrics import Accuracy
47 # self.train_accuracy = Accuracy(task="multiclass", num_classes=10)
48 # self.val_accuracy = Accuracy(task="multiclass", num_classes=10)
50 def forward(self, x):
51 """
52 Forward pass of the model.
54 Args:
55 x: Input tensor
57 Returns:
58 Model output
59 """
60 return self.model(x)
62 def training_step(self, batch, batch_idx):
63 """
64 Training step (called for each training batch).
66 Args:
67 batch: Current batch of data
68 batch_idx: Index of the current batch
70 Returns:
71 Loss tensor
72 """
73 x, y = batch
75 # Forward pass
76 logits = self(x)
77 loss = F.cross_entropy(logits, y)
79 # Calculate accuracy (optional)
80 preds = torch.argmax(logits, dim=1)
81 acc = (preds == y).float().mean()
83 # Log metrics
84 self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, batch_size=y.size(0), sync_dist=True)
85 self.log("train/acc", acc, on_step=True, on_epoch=True, batch_size=y.size(0), sync_dist=True)
86 self.log("learning_rate", self.optimizers().param_groups[0]["lr"])
88 return loss
90 def validation_step(self, batch, batch_idx):
91 """
92 Validation step (called for each validation batch).
94 Args:
95 batch: Current batch of data
96 batch_idx: Index of the current batch
97 """
98 x, y = batch
100 # Forward pass
101 logits = self(x)
102 loss = F.cross_entropy(logits, y)
104 # Calculate accuracy
105 preds = torch.argmax(logits, dim=1)
106 acc = (preds == y).float().mean()
108 # Log metrics (automatically aggregated across batches)
109 self.log("val/loss", loss, on_epoch=True, prog_bar=True, batch_size=y.size(0), sync_dist=True)
110 self.log("val/acc", acc, on_epoch=True, prog_bar=True, batch_size=y.size(0), sync_dist=True)
112 def test_step(self, batch, batch_idx):
113 """
114 Test step (called for each test batch).
116 Args:
117 batch: Current batch of data
118 batch_idx: Index of the current batch
119 """
120 x, y = batch
122 # Forward pass
123 logits = self(x)
124 loss = F.cross_entropy(logits, y)
126 # Calculate accuracy
127 preds = torch.argmax(logits, dim=1)
128 acc = (preds == y).float().mean()
130 # Log metrics
131 self.log("test/loss", loss, on_epoch=True, batch_size=y.size(0), sync_dist=True)
132 self.log("test/acc", acc, on_epoch=True, batch_size=y.size(0), sync_dist=True)
134 def predict_step(self, batch, batch_idx, dataloader_idx=0):
135 """
136 Prediction step (called for each prediction batch).
138 Args:
139 batch: Current batch of data
140 batch_idx: Index of the current batch
141 dataloader_idx: Index of the dataloader (if multiple)
143 Returns:
144 Predictions
145 """
146 x, y = batch
147 logits = self(x)
148 preds = torch.argmax(logits, dim=1)
149 return preds
151 def configure_optimizers(self):
152 """
153 Configure optimizers and learning rate schedulers.
155 Returns:
156 Optimizer and scheduler configuration
157 """
158 # Define optimizer
159 optimizer = Adam(
160 self.parameters(),
161 lr=self.hparams.learning_rate,
162 weight_decay=1e-5,
163 )
165 # Define scheduler. PyTorch removed the schedulers' `verbose` argument;
166 # use LearningRateMonitor or the logged learning_rate metric instead.
167 scheduler = ReduceLROnPlateau(
168 optimizer,
169 mode="min",
170 factor=0.5,
171 patience=5,
172 )
174 # Return configuration
175 return {
176 "optimizer": optimizer,
177 "lr_scheduler": {
178 "scheduler": scheduler,
179 "monitor": "val/loss",
180 "interval": "epoch",
181 "frequency": 1,
182 },
183 }
185 # Optional: Add custom methods for model-specific logic
187 def on_train_epoch_end(self):
188 """Called at the end of each training epoch."""
189 # Example: Log custom metrics
190 pass
192 def on_validation_epoch_end(self):
193 """Called at the end of each validation epoch."""
194 # Example: Compute epoch-level metrics
195 pass
198# Example usage
199if __name__ == "__main__":
200 # Create model
201 model = TemplateLightningModule(
202 learning_rate=0.001,
203 hidden_dim=256,
204 dropout=0.1,
205 )
207 # Create trainer
208 trainer = L.Trainer(
209 max_epochs=10,
210 accelerator="auto",
211 devices=1,
212 logger=False,
213 )
215 # Train (you need to provide train_dataloader and val_dataloader)
216 # trainer.fit(model, train_dataloader, val_dataloader)
218 n_parameters = sum(p.numel() for p in model.parameters())
219 print(f"Model created with {n_parameters:,} parameters")
220 print(f"Hyperparameters: {model.hparams}")