alexnet.py

# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later

import lightning.pytorch as pl
import torch
import torch.nn as nn
import torchvision.models as models

from .normalizer import TorchVisionNormalizer


class Alexnet(pl.LightningModule):
    """Alexnet module.

    Note: only usable with a normalized dataset
    """

    def __init__(
        self,
        criterion,
        criterion_valid,
        optimizer,
        optimizer_configs,
        pretrained=False,
    ):
        super().__init__()

        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])

        self.name = "AlexNet"

        # Load pretrained model
        weights = (
            None if pretrained is False else models.AlexNet_Weights.DEFAULT
        )
        self.model_ft = models.alexnet(weights=weights)

        self.normalizer = TorchVisionNormalizer(nb_channels=1)

        # Adapt output features
        self.model_ft.classifier[4] = nn.Linear(4096, 512)
        self.model_ft.classifier[6] = nn.Linear(512, 1)

    def forward(self, x):
        x = self.normalizer(x)
        x = self.model_ft(x)

        return x

    def training_step(self, batch, batch_idx):
        images = batch[1]
        labels = batch[2]

        # Increase label dimension if too low
        # Allows single and multiclass usage
        if labels.ndim == 1:
            labels = torch.reshape(labels, (labels.shape[0], 1))

        # Forward pass on the network
        outputs = self(images)

        # Manually move criterion to selected device, since not part of the model.
        self.hparams.criterion = self.hparams.criterion.to(self.device)
        training_loss = self.hparams.criterion(outputs, labels.float())

        return {"loss": training_loss}

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        images = batch[1]
        labels = batch[2]

        # Increase label dimension if too low
        # Allows single and multiclass usage
        if labels.ndim == 1:
            labels = torch.reshape(labels, (labels.shape[0], 1))

        # data forwarding on the existing network
        outputs = self(images)

        # Manually move criterion to selected device, since not part of the model.
        self.hparams.criterion_valid = self.hparams.criterion_valid.to(
            self.device
        )
        validation_loss = self.hparams.criterion_valid(outputs, labels.float())

        if dataloader_idx == 0:
            return {"validation_loss": validation_loss}
        else:
            return {f"extra_validation_loss_{dataloader_idx}": validation_loss}

    def predict_step(self, batch, batch_idx, dataloader_idx=0, grad_cams=False):
        names = batch[0]
        images = batch[1]

        outputs = self(images)
        probabilities = torch.sigmoid(outputs)

        # necessary check for HED architecture that uses several outputs
        # for loss calculation instead of just the last concatfuse block
        if isinstance(outputs, list):
            outputs = outputs[-1]

        return names[0], torch.flatten(probabilities), torch.flatten(batch[2])

    def configure_optimizers(self):
        optimizer = getattr(torch.optim, self.hparams.optimizer)(
            self.parameters(), **self.hparams.optimizer_configs
        )

        return optimizer