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, )