src/models/focus_litenn_module.py

import math
from typing import Any, List

import torch
import torch.nn.functional as F
from torch import nn
from pytorch_lightning import LightningModule
from torchmetrics import MaxMetric, MeanAbsoluteError, MinMetric
from torchmetrics.classification.accuracy import Accuracy
import torchvision.models as models


class FocusLiteNN(nn.Module):
    """
    https://github.com/icbcbicc/FocusLiteNN
    This class is licenced undeer:
    # The Prosperity Public License 3.0.0
    Contributor: Zhongling Wang, Mahdi S. Hosseini, Adyn Miles
    Source Code: https://github.com/icbcbicc/FocusLiteNN

    ## Purpose:
    This license allows you to use and share this software for noncommercial purposes for free and to try this software for commercial purposes for thirty days.

    ## Agreement
    In order to receive this license, you have to agree to its rules.  Those rules are both obligations under that agreement and conditions to your license.  Don't do anything with this software that triggers a rule you can't or won't follow.

    ## Notices
    Make sure everyone who gets a copy of any part of this software from you, with or without changes, also gets the text of this license and the contributor and source code lines above.

    ## Commercial Trial
    Limit your use of this software for commercial purposes to a thirty-day trial period.  If you use this software for work, your company gets one trial period for all personnel, not one trial per person.

    ## Contributions Back
    Developing feedback, changes, or additions that you contribute back to the contributor on the terms of a standardized public software license such as [the Blue Oak Model License 1.0.0](https://blueoakcouncil.org/license/1.0.0), [the Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0.html), [the MIT license](https://spdx.org/licenses/MIT.html), or [the two-clause BSD license](https://spdx.org/licenses/BSD-2-Clause.html) doesn't count as use for a commercial purpose.

    ## Personal Uses
    Personal use for research, experiment, and testing for the benefit of public knowledge, personal study, private entertainment, hobby projects, amateur pursuits, or religious observance, without any anticipated commercial application, doesn't count as use for a commercial purpose.

    ## Noncommercial Organizations
    Use by any charitable organization, educational institution, public research organization, public safety or health organization, environmental protection organization, or government institution doesn't count as use for a commercial purpose regardless of the source of funding or obligations resulting from the funding.

    ## Defense
    Don't make any legal claim against anyone accusing this software, with or without changes, alone or with other technology, of infringing any patent.

    ## Copyright
    The contributor licenses you to do everything with this software that would otherwise infringe their copyright in it.

    ## Patent
    The contributor licenses you to do everything with this software that would otherwise infringe any patents they can license or become able to license.

    ## Reliability
    The contributor can't revoke this license.

    ## Excuse
    You're excused for unknowingly breaking [Notices](#notices) if you take all practical steps to comply within thirty days of learning you broke the rule.

    ## No Liability
    ***As far as the law allows, this software comes as is, without any warranty or condition, and the contributor won't be liable to anyone for any damages related to this software or this license, under any kind of legal claim.***

    In practice, we found MIN contributes VERY LITTLE to the performance. To achieve
    extreme simplicity, different from Equation 1 in the paper, we only use MAX as the
    nonlinear function. All experimental results in the paper are reported based on the
    model using only MAX.
    The model (class FocusLiteNNMinMax) using weighted MAX and MIN as the nonlinear function
    (Equation 1 in paper) has indistinguishable (slightly better) performance compared to the
    model using only MAX (class FocusLiteNN).
    """

    def __init__(self, num_channel=3):
        super(FocusLiteNN, self).__init__()
        self.num_channel = num_channel
        self.conv = nn.Conv2d(3, self.num_channel, 7, stride=5, padding=1)  # 47x47
        self.maxpool = nn.MaxPool2d(kernel_size=47)
        if self.num_channel > 1:
            self.fc = nn.Conv2d(self.num_channel, 1, 1, stride=1, padding=0)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))

    def forward(self, x):
        batch_size = x.size()[0]

        x = self.conv(x)
        x = -self.maxpool(-x)  # minpooling
        if self.num_channel > 1:
            x = self.fc(x)
        x = x.view(batch_size, -1)

        return x


"""
All code below is licenced under the licence of the whole repository
"""


class FocusLiteNNLitModule(LightningModule):
    def __init__(
        self,
        lr: float = 0.001,
        num_channel: int = 3,
        pre_trained: bool = False,
        weight_decay: float = 0.0005,
    ):
        super().__init__()

        # this line allows to access init params with 'self.hparams' attribute
        # it also ensures init params will be stored in ckpt
        self.save_hyperparameters(logger=False)

        # loss function
        self.criterion = torch.nn.MSELoss()

        # use separate metric instance for train, val and test step
        # to ensure a proper reduction over the epoch
        self.train_mae = MeanAbsoluteError()
        self.val_mae = MeanAbsoluteError()
        self.test_mae = MeanAbsoluteError()

        # for logging best so far validation accuracy
        self.val_mae_best = MinMetric()

        self.num_channel = num_channel
        self.model = FocusLiteNN(num_channel)

        self.model.maxpool = nn.AdaptiveMaxPool2d(1)

    def forward(self, x):
        x = self.model(x)
        return x

    def step(self, batch: Any):
        x = batch["image"]
        y = batch["focus_height"]
        logits = self.forward(x)
        loss = self.criterion(logits, y.unsqueeze(1))
        preds = torch.squeeze(logits)
        return loss, preds, y

    def training_step(self, batch: Any, batch_idx: int):
        loss, preds, targets = self.step(batch)

        # log train metrics
        mae = self.train_mae(preds, targets)
        self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=False)
        self.log("train/mae", mae, on_step=False, on_epoch=True, prog_bar=True)

        # we can return here dict with any tensors
        # and then read it in some callback or in `training_epoch_end()`` below
        # remember to always return loss from `training_step()` or else
        # backpropagation will fail!
        return {"loss": loss, "preds": preds, "targets": targets}

    def training_epoch_end(self, outputs: List[Any]):
        # `outputs` is a list of dicts returned from `training_step()`
        pass

    def validation_step(self, batch: Any, batch_idx: int):
        loss, preds, targets = self.step(batch)

        # log val metrics
        mae = self.val_mae(preds, targets)
        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=False)
        self.log("val/mae", mae, on_step=False, on_epoch=True, prog_bar=True)

        return {"loss": loss, "preds": preds, "targets": targets}

    def validation_epoch_end(self, outputs: List[Any]):
        mae = self.val_mae.compute()  # get val accuracy from current epoch
        self.val_mae_best.update(mae)
        self.log(
            "val/mae_best", self.val_mae_best.compute(), on_epoch=True, prog_bar=True
        )

    def test_step(self, batch: Any, batch_idx: int):
        loss, preds, targets = self.step(batch)

        # log test metrics
        mae = self.test_mae(preds, targets)
        self.log("test/loss", loss, on_step=False, on_epoch=True)
        self.log("test/mae", mae, on_step=False, on_epoch=True)

    def test_epoch_end(self, outputs: List[Any]):
        print(outputs)
        pass

    def on_epoch_end(self):
        # reset metrics at the end of every epoch
        self.train_mae.reset()
        self.test_mae.reset()
        self.val_mae.reset()

    def configure_optimizers(self):
        """Choose what optimizers and learning-rate schedulers.

        Normally you'd need one. But in the case of GANs or similar you might
        have multiple.

        See examples here:
            https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers
        """
        return torch.optim.Adam(
            params=self.parameters(),
            lr=self.hparams.lr,
            weight_decay=self.hparams.weight_decay,
        )