ptame.models.cv_module

  1from collections.abc import Callable, Iterator
  2from functools import partial
  3from typing import Any
  4
  5import torch
  6from lightning import LightningModule
  7from torch import Tensor
  8from torchmetrics import MaxMetric, MeanMetric
  9from torchmetrics.classification.accuracy import Accuracy
 10
 11
 12class CVModule(LightningModule):
 13    """`LightningModule` for PAMELA.
 14
 15    A `LightningModule` implements 8 key methods:
 16
 17    ```python
 18    def __init__(self):
 19    # Define initialization code here.
 20
 21    def setup(self, stage):
 22    # Things to setup before each stage, 'fit', 'validate', 'test', 'predict'.
 23    # This hook is called on every process when using DDP.
 24
 25    def training_step(self, batch, batch_idx):
 26    # The complete training step.
 27
 28    def validation_step(self, batch, batch_idx):
 29    # The complete validation step.
 30
 31    def test_step(self, batch, batch_idx):
 32    # The complete test step.
 33
 34    def predict_step(self, batch, batch_idx):
 35    # The complete predict step.
 36
 37    def configure_optimizers(self):
 38    # Define and configure optimizers and LR schedulers.
 39    ```
 40
 41    Docs:
 42        https://lightning.ai/docs/pytorch/latest/common/lightning_module.html
 43    """
 44
 45    def __init__(
 46        self,
 47        net: torch.nn.Module,
 48        loss: Callable[[Tensor], Tensor],
 49        optimizer: Callable[
 50            [Iterator[torch.nn.Parameter]], torch.optim.Optimizer
 51        ],
 52        scheduler: dict[str, Any] | None,
 53        teacher: Callable[[Tensor], Tensor] | None = None,
 54        terminate_on_nan: bool = True,
 55        compile: bool = False,
 56        **kwargs,
 57    ) -> None:
 58        """Initialize a `PAMELALitModule`.
 59
 60        :param net: The model to train.
 61        :param optimizer: The optimizer to use for training.
 62        :param scheduler: The learning rate scheduler to use for training,
 63            together with the options (how often it should update etc.).
 64        :param val_measures: The explainability specific validation measures to
 65            use.
 66        :param test_measures: " test ".
 67        :param compile: Whether to compile the model before training.
 68        """
 69        super().__init__()
 70
 71        # this line allows to access init params with 'self.hparams' attribute
 72        # also ensures init params will be stored in ckpt
 73        self.save_hyperparameters(
 74            ignore=["net", "loss", "optimizer", "scheduler", "teacher"],
 75            logger=False,
 76        )
 77
 78        self.net = net
 79        self.teacher = None
 80        if teacher is not None:
 81            self.teacher = teacher
 82            loss = self.kl_loss
 83            self.teacher.eval()
 84            self.teacher.requires_grad_(False)
 85
 86        # cut off initialization for simple restore
 87        if loss is None:
 88            return
 89
 90        # loss function
 91        self.criterion = loss
 92
 93        # optimizer and scheduler
 94        self.optimizer_cfg = optimizer
 95        self.scheduler_cfg = scheduler
 96
 97        # metric objects are initialized in `setup` method
 98        self.train_acc = None
 99        self.val_acc = None
100        self.test_acc = None
101        # for averaging loss across batches
102        self.train_loss = MeanMetric()
103        self.val_loss = MeanMetric()
104        self.test_loss = MeanMetric()
105        # metric object for calculating and averaging ADIC or ROAD or across batches
106
107        # for tracking best so far validation accuracy
108        self.val_acc_best = MaxMetric()
109
110    @staticmethod
111    def mse_loss(student_logits, teacher_logits):
112        return torch.nn.functional.mse_loss(student_logits, teacher_logits)
113
114    @staticmethod
115    def kl_loss(student_logits, teacher_logits, temperature=5.0):
116        p_t = torch.nn.functional.softmax(teacher_logits / temperature, dim=-1)
117        p_s = torch.nn.functional.softmax(student_logits / temperature, dim=-1)
118        return torch.nn.functional.kl_div(
119            p_s.log(), p_t, reduction="batchmean"
120        )
121
122    def convert_fc(self) -> None:
123        """Convert the last layer of the model to a fully connected layer."""
124        if hasattr(self.net, "fc"):
125            in_features = self.net.fc.in_features
126            out_features = (
127                self.trainer.datamodule.num_classes
128                if self.teacher is None
129                else self.teacher.fc.out_features
130            )
131            self.classes = out_features
132            if self.net.fc.out_features == out_features:
133                return
134            self.net.fc = torch.nn.Linear(in_features, out_features)
135            # initialize weights
136            torch.nn.init.xavier_uniform_(self.net.fc.weight)
137            if self.net.fc.bias is not None:
138                torch.nn.init.zeros_(self.net.fc.bias)
139
140    def forward(self, x: torch.Tensor) -> torch.Tensor:
141        """Perform a forward pass through the model `self.net`.
142
143        :param x: A tensor of images.
144        :return: A tensor of logits.
145        """
146        if self.hparams.compile:
147            return self.compiled_net(x)
148        return self.net(x)
149
150    def on_train_start(self) -> None:
151        """Lightning hook that is called when training begins."""
152        # by default lightning executes validation step sanity checks before training starts,
153        # so it's worth to make sure validation metrics don't store results from these checks
154        self.val_loss.reset()
155        self.val_acc.reset()
156        self.val_acc_best.reset()
157
158    def model_step(
159        self, batch: tuple[torch.Tensor, torch.Tensor]
160    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
161        """Perform a single model step on a batch of data.
162
163        :param batch: A batch of data (a tuple) containing the input tensor of images and target labels.
164
165        :return: A tuple containing (in order):
166            - A tensor of losses.
167            - A tensor of predictions.
168            - A tensor of target labels.
169        """
170        x, y = batch
171        target = y
172        if self.teacher is not None:
173            if self.hparams.compile:
174                y = self.compiled_teacher(x)
175            else:
176                y = self.teacher(x)
177            target = y.argmax(dim=1)
178        out = self.forward(x)
179        loss = self.criterion(out, y)
180        return loss, out, target
181
182    def training_step(
183        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int
184    ) -> torch.Tensor:
185        """Perform a single training step on a batch of data from the training
186        set.
187
188        :param batch: A batch of data (a tuple) containing the input tensor of
189            images and target labels.
190        :param batch_idx: The index of the current batch.
191        :return: A tensor of losses between model predictions and targets.
192        """
193        loss, out, target = self.model_step(batch)
194        preds = out.argmax(dim=1)
195        # update and log metrics
196        self.train_loss(loss)  # compute metric
197        self.log("train/loss", self.train_loss, prog_bar=True)
198        self.train_acc(preds, target)
199        self.log("train/acc", self.train_acc, prog_bar=True)
200
201        # return loss or backpropagation will fail
202        if self.hparams.terminate_on_nan and loss.isnan().any():
203            raise ValueError("NaN detected in loss!")
204        return loss
205
206    def validation_step(
207        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int
208    ) -> None:
209        """Perform a single validation step on a batch of data from the
210        validation set.
211
212        :param batch: A batch of data (a tuple) containing the input tensor of
213            images and target labels.
214        :param batch_idx: The index of the current batch.
215        """
216        loss, out, target = self.model_step(batch)
217        preds = out.argmax(dim=1)
218
219        # update and log metrics
220        self.val_loss(loss)
221        self.log("val/loss", self.val_loss)
222        self.val_acc(preds, target)
223        self.log("val/acc", self.val_acc)
224
225    def on_validation_epoch_end(self) -> None:
226        "Lightning hook that is called when a validation epoch ends."
227        acc = self.val_acc.compute()  # get current val acc
228        self.val_acc_best(acc)  # update best so far val acc
229        # log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
230        # otherwise metric would be reset by lightning after each epoch
231        self.log(
232            "val/acc_best",
233            self.val_acc_best.compute(),
234            sync_dist=True,
235            prog_bar=True,
236        )
237
238    def test_step(
239        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int
240    ) -> None:
241        """Perform a single test step on a batch of data from the test set.
242
243        :param batch: A batch of data (a tuple) containing the input tensor of
244            images and target labels.
245        :param batch_idx: The index of the current batch.
246        """
247        loss, out, target = self.model_step(batch)
248        preds = out.argmax(dim=1)
249
250        # update and log metrics
251        self.test_loss(loss)
252        self.log("test/loss", self.test_loss)
253        self.test_acc(preds, target)
254        self.log("test/acc", self.test_acc)
255
256    def setup(self, stage: str) -> None:
257        """Lightning hook that is called at the beginning of fit (train +
258        validate), validate, test, or predict.
259
260        This is a good hook when you need to build models dynamically or adjust
261        something about them. This hook is called on every process when using
262        DDP.
263
264        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
265        """
266        self.convert_fc()
267        if self.hparams.compile:
268            self.compiled_net = torch.compile(self.net)
269            if self.teacher is not None:
270                self.compiled_teacher = torch.compile(self.teacher)
271        # metric objects for calculating and averaging accuracy across batches
272        self.train_acc = Accuracy(task="multiclass", num_classes=self.classes)
273        self.val_acc = Accuracy(task="multiclass", num_classes=self.classes)
274        self.test_acc = Accuracy(task="multiclass", num_classes=self.classes)
275
276    def teardown(self, stage: str) -> None:
277        """Lightning hook that is called at the end of fit (train + validate),
278        validate, test, or predict.
279
280        This is a good hook when you need to clean something up after the run.
281
282        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
283        """
284        if self.hparams.compile:
285            del self.compiled_net
286            if self.teacher:
287                del self.compiled_teacher
288
289    def configure_optimizers(self) -> dict[str, Any]:
290        """Choose what optimizers and learning-rate schedulers to use in your
291        optimization. Normally you'd need one. But in the case of GANs or
292        similar you might have multiple.
293
294        Examples:
295            https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
296
297        :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
298        """
299        optimizer = self.optimizer_cfg(params=self.trainer.model.parameters())
300        if (scheduler_info := self.scheduler_cfg) is not None:
301            scheduler = scheduler_info["scheduler"]
302            if scheduler_info["name"] == "OneCycleLR":
303                scheduler = partial(
304                    scheduler,
305                    total_steps=self.trainer.estimated_stepping_batches,
306                )
307            scheduler = scheduler(optimizer=optimizer)
308            scheduler_params = scheduler_info["params"]
309            return {
310                "optimizer": optimizer,
311                "lr_scheduler": {
312                    "scheduler": scheduler,
313                    "interval": scheduler_params["interval"],
314                    "frequency": scheduler_params["frequency"],
315                    "name": scheduler_params["name"],
316                },
317            }
318        return {"optimizer": optimizer}
319
320    def on_save_checkpoint(self, checkpoint):
321        # pop the keys you are not interested by
322        sd = checkpoint["state_dict"]
323        names = list(sd.keys())
324        for name in names:
325            if any(x in name for x in ["_orig_mod", "fc", "teacher"]):
326                sd.pop(name)
327
328
329if __name__ == "__main__":
330    _ = CVModule(None, None, None, None, None, None)