Add simple fully convolutional network

configs/experiment/focusConvMSE_150.yaml

@@ -0,0 +1,39 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+defaults:
+  - override /datamodule: focus150.yaml
+  - override /model: focusConv_150.yaml
+  - override /callbacks: default.yaml
+  - override /logger: many_loggers
+  - override /trainer: default.yaml
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+# name of the run determines folder name in logs
+name: "focusConvMSE_150"
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 100
+
+model:
+  image_size: 150
+  pool_size: 2
+  conv1_size: 5
+  conv1_channels: 6
+  conv2_size: 5
+  conv2_channels: 16
+  lin1_size: 100
+  lin2_size: 80
+  output_size: 1
+  lr: 0.001
+  weight_decay: 0.0005
+
+datamodule:
+  batch_size: 128
+  

configs/model/focusConv_150.yaml

@@ -0,0 +1,13 @@
+_target_: src.models.focus_conv_module.FocusConvLitModule
+
+image_size: 150
+pool_size: 2
+conv1_size: 5
+conv1_channels: 6
+conv2_size: 5
+conv2_channels: 16
+lin1_size: 100
+lin2_size: 80
+output_size: 1
+lr: 0.001
+weight_decay: 0.0005

src/models/focus_conv_module.py

@@ -9,54 +9,42 @@ from torchmetrics.classification.accuracy import Accuracy
 
 
 class SimpleConvNet(nn.Module):
-    def __init__(self):
+    def __init__(self, hparams):
         super().__init__()
-        self.conv1 = nn.Conv2d(3, 6, 5)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(6, 16, 5)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(6, 16, 5)
-        self.fc1 = nn.Linear(16 * 5 * 5, 120)
-        self.fc2 = nn.Linear(120, 84)
-        self.fc3 = nn.Linear(84, 10)
 
-    def forward(self, x):
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = torch.flatten(x, 1)  # flatten all dimensions except batch
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
-        x = self.fc3(x)
-        return x
+        pool_size = hparams["pool_size"]  # 2
+        conv1_size = hparams["conv1_size"]  # 5
+        conv1_out = hparams["conv1_channels"]  # 6
+        conv2_size = hparams["conv1_channels"]  # 5
+        conv2_out = hparams["conv2_channels"]  # 16
+        size_img = hparams["image_size"]  # 150
 
+        lin1_size = hparams["lin1_size"]  # 100
+        lin2_size = hparams["lin2_size"]  # 80
+        output_size = hparams["output_size"]  # 1
 
-class SimpleDenseNet(nn.Module):
-    def __init__(self, hparams: dict):
-        super().__init__()
+        size_img -= conv1_size - 1
+        size_img = int((size_img) / pool_size)
+        size_img -= conv2_size - 1
+        size_img = int(size_img / pool_size)
 
         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"]),
+            nn.Conv2d(3, conv1_out, conv1_size),
+            nn.MaxPool2d(pool_size, pool_size),
+            nn.Conv2d(conv1_out, conv2_out, conv2_size),
+            nn.MaxPool2d(pool_size, pool_size),
+            nn.Flatten(),
+            nn.Linear(conv2_out * size_img * size_img, lin1_size),
+            nn.Linear(lin1_size, lin2_size),
+            nn.Linear(lin2_size, 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)
+        x = self.model(x)
+        return x
 
 
-class FocusLitModule(LightningModule):
+class FocusConvLitModule(LightningModule):
     """
     Example of LightningModule for MNIST classification.
 
@@ -73,10 +61,14 @@ class FocusLitModule(LightningModule):
 
     def __init__(
         self,
-        input_size: int = 75 * 75 * 3,
-        lin1_size: int = 256,
-        lin2_size: int = 256,
-        lin3_size: int = 256,
+        image_size: int = 150,
+        pool_size: int = 2,
+        conv1_size: int = 5,
+        conv1_channels: int = 6,
+        conv2_size: int = 5,
+        conv2_channels: int = 16,
+        lin1_size: int = 100,
+        lin2_size: int = 80,
         output_size: int = 1,
         lr: float = 0.001,
         weight_decay: float = 0.0005,
@@ -87,10 +79,10 @@ class FocusLitModule(LightningModule):
         # it also ensures init params will be stored in ckpt
         self.save_hyperparameters(logger=False)
 
-        self.model = SimpleDenseNet(hparams=self.hparams)
+        self.model = SimpleConvNet(hparams=self.hparams)
 
         # loss function
-        self.criterion = torch.nn.L1Loss()
+        self.criterion = torch.nn.MSELoss()
 
         # use separate metric instance for train, val and test step
         # to ensure a proper reduction over the epoch
@@ -108,7 +100,7 @@ class FocusLitModule(LightningModule):
         x = batch["image"]
         y = batch["focus_value"]
         logits = self.forward(x)
-        loss = self.criterion(logits, y)
+        loss = self.criterion(logits, y.unsqueeze(1))
         preds = torch.squeeze(logits)
         return loss, preds, y
 
@@ -122,7 +114,8 @@ class FocusLitModule(LightningModule):
 
         # 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!
+        # 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]):
@@ -169,7 +162,8 @@ class FocusLitModule(LightningModule):
     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.
+        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