-
Notifications
You must be signed in to change notification settings - Fork 1
/
mnist_module.py
144 lines (115 loc) · 5 KB
/
mnist_module.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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
from typing import Any
import torch
from lightning import LightningModule
from torchmetrics import MaxMetric, MeanMetric
from torchmetrics.classification.accuracy import Accuracy
class MNISTLitModule(LightningModule):
"""Example of LightningModule for MNIST classification.
A LightningModule organizes your PyTorch code into 6 sections:
- Initialization (__init__)
- Train Loop (training_step)
- Validation loop (validation_step)
- Test loop (test_step)
- Prediction Loop (predict_step)
- Optimizers and LR Schedulers (configure_optimizers)
Docs:
https://lightning.ai/docs/pytorch/latest/common/lightning_module.html
"""
def __init__(
self,
net: torch.nn.Module,
optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler,
huggingface_model: str = None, # only for compatibility
):
super().__init__()
# 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.net = net
# loss function
self.criterion = torch.nn.CrossEntropyLoss()
# metric objects for calculating and averaging accuracy across batches
self.train_acc = Accuracy(task="multiclass", num_classes=10)
self.val_acc = Accuracy(task="multiclass", num_classes=10)
self.test_acc = Accuracy(task="multiclass", num_classes=10)
# for averaging loss across batches
self.train_loss = MeanMetric()
self.val_loss = MeanMetric()
self.test_loss = MeanMetric()
# for tracking best so far validation accuracy
self.val_acc_best = MaxMetric()
def forward(self, x: torch.Tensor):
return self.net(x)
def on_train_start(self):
# by default lightning executes validation step sanity checks before training starts,
# so it's worth to make sure validation metrics don't store results from these checks
self.val_loss.reset()
self.val_acc.reset()
self.val_acc_best.reset()
def model_step(self, batch: Any):
x, y = batch
logits = self.forward(x)
loss = self.criterion(logits, y)
preds = torch.argmax(logits, dim=1)
return loss, preds, y
def training_step(self, batch: Any, batch_idx: int):
loss, preds, targets = self.model_step(batch)
# update and log metrics
self.train_loss(loss)
self.train_acc(preds, targets)
self.log(
"train/loss", self.train_loss, on_step=False, on_epoch=True, prog_bar=True
)
self.log(
"train/acc", self.train_acc, on_step=False, on_epoch=True, prog_bar=True
)
# return loss or backpropagation will fail
return loss
def on_train_epoch_end(self):
pass
def validation_step(self, batch: Any, batch_idx: int):
loss, preds, targets = self.model_step(batch)
# update and log metrics
self.val_loss(loss)
self.val_acc(preds, targets)
self.log("val/loss", self.val_loss, on_step=False, on_epoch=True, prog_bar=True)
self.log("val/acc", self.val_acc, on_step=False, on_epoch=True, prog_bar=True)
def on_validation_epoch_end(self):
acc = self.val_acc.compute() # get current val acc
self.val_acc_best(acc) # update best so far val acc
# log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
# otherwise metric would be reset by lightning after each epoch
self.log("val/acc_best", self.val_acc_best.compute(), prog_bar=True)
def test_step(self, batch: Any, batch_idx: int):
loss, preds, targets = self.model_step(batch)
# update and log metrics
self.test_loss(loss)
self.test_acc(preds, targets)
self.log(
"test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True
)
self.log("test/acc", self.test_acc, on_step=False, on_epoch=True, prog_bar=True)
def on_test_epoch_end(self):
pass
def configure_optimizers(self):
"""Choose what optimizers and learning-rate schedulers to use in your optimization.
Normally you'd need one. But in the case of GANs or similar you might have multiple.
Examples:
https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
"""
optimizer = self.hparams.optimizer(params=self.parameters())
if self.hparams.scheduler is not None:
scheduler = self.hparams.scheduler(optimizer=optimizer)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val/loss",
"interval": "epoch",
"frequency": 1,
},
}
return {"optimizer": optimizer}
if __name__ == "__main__":
_ = MNISTLitModule(None, None, None)