implement simple model

models/notebooks/1.0-hfk-datamodules-exploration.ipynb

@@ -393,7 +393,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -402,7 +402,7 @@
        "64"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -413,6 +413,53 @@
     "\n",
     "len(data[\"focus_value\"])"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/hku/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py:96: UserWarning: Using a target size (torch.Size([64])) that is different to the input size (torch.Size([64, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
+      "  return F.l1_loss(input, target, reduction=self.reduction)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "(tensor(2.5787, grad_fn=<L1LossBackward0>),\n",
+       " tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
+       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
+       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),\n",
+       " tensor([-1.2805, -0.0943, -2.3645,  0.8542, -0.8047, -6.0020,  0.0000, -4.3352,\n",
+       "         -1.8066, -2.7189, -6.4697, -3.2557, -4.2778, -5.0264, -3.4891,  0.0000,\n",
+       "         -1.7181, -2.7314,  0.3324, -0.0943, -0.8991,  0.0000, -4.4178,  1.9723,\n",
+       "         -3.0026, -5.5685,  3.8374,  3.8625, -0.4125, -4.1936, -1.5781, -1.6393,\n",
+       "         -2.9583, -5.4933, -1.7807, -3.3135, -5.3423, -0.7978, -5.3971, -4.9412,\n",
+       "          0.0000, -4.4128, -5.7744, -5.2755, -1.0996, -5.7482,  0.0000, -0.1737,\n",
+       "         -3.5851, -6.1429, -6.3642, -3.9653, -0.2081, -0.9539, -0.4159, -0.5388,\n",
+       "         -1.3643, -4.4441, -1.5161,  0.6395, -5.4710, -2.6482,  0.0000, -2.6257],\n",
+       "        dtype=torch.float64))"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import types\n",
+    "import importlib.machinery\n",
+    "focus_module = SourceFileLoader('focus_module', '../src/models/focus_module.py').load_module()\n",
+    "from focus_module import FocusLitModule\n",
+    "\n",
+    "model = FocusLitModule()\n",
+    "\n",
+    "model.step(data)"
+   ]
   }
  ],
  "metadata": {

models/requirements.txt

@@ -6,6 +6,7 @@ torchmetrics>=0.7.0
 
 # --------- data and model dependencies --------- #
 scikit-image
+pandas
 
 # --------- hydra --------- #
 hydra-core>=1.1.0

models/src/datamodules/focus_datamodule.py

@@ -51,7 +51,7 @@ class FocusDataSet(Dataset):
         sample = {"image": image, "focus_value": focus_value}
 
         if self.transform:
-            sample = self.transform(sample)
+            sample["image"] = self.transform(sample["image"])
 
         return sample
 
@@ -76,7 +76,9 @@ class FocusDataModule(LightningDataModule):
         self.save_hyperparameters(logger=False)
 
         # data transformations
-        self.transforms = transforms.Compose([])
+        self.transforms = transforms.Compose(
+            [transforms.ToTensor(), transforms.ConvertImageDtype(torch.float)]
+        )
 
         self.data_train: Optional[Dataset] = None
         self.data_val: Optional[Dataset] = None

models/src/models/focus_module.py

@@ -0,0 +1,157 @@
+from typing import Any, List
+
+import torch
+from torch import nn
+from pytorch_lightning import LightningModule
+from torchmetrics import MaxMetric, MeanAbsoluteError, MinMetric
+from torchmetrics.classification.accuracy import Accuracy
+
+
+class SimpleDenseNet(nn.Module):
+    def __init__(self, hparams: dict):
+        super().__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(hparams["input_size"], hparams["lin1_size"]),
+            nn.BatchNorm1d(hparams["lin1_size"]),
+            nn.ReLU(),
+            nn.Linear(hparams["lin1_size"], hparams["lin2_size"]),
+            nn.BatchNorm1d(hparams["lin2_size"]),
+            nn.ReLU(),
+            nn.Linear(hparams["lin2_size"], hparams["lin3_size"]),
+            nn.BatchNorm1d(hparams["lin3_size"]),
+            nn.ReLU(),
+            nn.Linear(hparams["lin3_size"], hparams["output_size"]),
+        )
+
+    def forward(self, x):
+        batch_size, channels, width, height = x.size()
+
+        # (batch, 1, width, height) -> (batch, 1*width*height)
+        x = x.view(batch_size, -1)
+
+        return self.model(x)
+
+
+class FocusLitModule(LightningModule):
+    """
+    Example of LightningModule for MNIST classification.
+
+    A LightningModule organizes your PyTorch code into 5 sections:
+        - Computations (init).
+        - Train loop (training_step)
+        - Validation loop (validation_step)
+        - Test loop (test_step)
+        - Optimizers (configure_optimizers)
+
+    Read the docs:
+        https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html
+    """
+
+    def __init__(
+        self,
+        input_size: int = 75 * 75 * 3,
+        lin1_size: int = 256,
+        lin2_size: int = 256,
+        lin3_size: int = 256,
+        output_size: int = 1,
+        lr: float = 0.001,
+        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)
+
+        self.model = SimpleDenseNet(hparams=self.hparams)
+
+        # loss function
+        self.criterion = torch.nn.L1Loss()
+
+        # use separate metric instance for train, val and test step
+        # to ensure a proper reduction over the epoch
+        self.train_acc = MeanAbsoluteError()
+        self.val_acc = MeanAbsoluteError()
+        self.test_acc = MeanAbsoluteError()
+
+        # for logging best so far validation accuracy
+        self.val_acc_best = MinMetric()
+
+    def forward(self, x: torch.Tensor):
+        return self.model(x)
+
+    def step(self, batch: Any):
+        x = batch["image"]
+        y = batch["focus_value"]
+        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.step(batch)
+
+        # log train metrics
+        acc = self.train_acc(preds, targets)
+        self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=False)
+        self.log("train/acc", acc, 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
+        acc = self.val_acc(preds, targets)
+        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=False)
+        self.log("val/acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+
+        return {"loss": loss, "preds": preds, "targets": targets}
+
+    def validation_epoch_end(self, outputs: List[Any]):
+        acc = self.val_acc.compute()  # get val accuracy from current epoch
+        self.val_acc_best.update(acc)
+        self.log(
+            "val/acc_best", self.val_acc_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
+        acc = self.test_acc(preds, targets)
+        self.log("test/loss", loss, on_step=False, on_epoch=True)
+        self.log("test/acc", acc, on_step=False, on_epoch=True)
+
+        return {"loss": loss, "preds": preds, "targets": targets}
+
+    def test_epoch_end(self, outputs: List[Any]):
+        pass
+
+    def on_epoch_end(self):
+        # reset metrics at the end of every epoch
+        self.train_acc.reset()
+        self.test_acc.reset()
+        self.val_acc.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,
+        )