rename default training into distance training and fix issues with usage as abstract training

mu_map/dataset/default.py

@@ -76,7 +76,7 @@ class MuMapDataset(Dataset):
         self.reconstructions = {}
         self.mu_maps = {}
 
-    def split_copy(self, split_name: str) -> MuMapDataset:
+    def split_copy(self, split_name: str):
         return MuMapDataset(
             dataset_dir=self.dir,
             csv_file=os.path.basename(self.csv_file),

mu_map/dataset/patches.py

@@ -45,7 +45,7 @@ class MuMapPatchDataset(MuMapDataset):
         self.patches = []
         self.generate_patches()
 
-    def split_copy(self, split_name: str) ->MuMapPatchDataset:
+    def split_copy(self, split_name: str):
         return MuMapPatchDataset(
             dataset_dir=self.dir,
             patches_per_image=self.patches_per_image,

mu_map/training/default.py → mu_map/training/distance.py

@@ -3,6 +3,7 @@ from typing import Dict
 
 import torch
 
+from mu_map.dataset.default import MuMapDataset
 from mu_map.logging import get_logger
 from mu_map.training.lib import TrainingParams, AbstractTraining
 from mu_map.training.loss import WeightedLoss
@@ -17,25 +18,25 @@ class Training(AbstractTraining):
         device: torch.device,
         snapshot_dir: str,
         snapshot_epoch: int,
-
         params: TrainingParams,
         loss_func: WeightedLoss,
+        logger,
     ):
-        super().__init__(epochs, dataset, batch_size, device, snapshot_dir, snapshot_epoch)
+        super().__init__(epochs, dataset, batch_size, device, snapshot_dir, snapshot_epoch, logger)
         self.training_params.append(params)
 
         self.loss_func = loss_func
         self.model = params.model
 
     def _train_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
-        outputs = self.model(inputs)
-        loss = self.loss_func(outputs, mu_maps)
+        mu_maps_predicted = self.model(recons)
+        loss = self.loss_func(mu_maps_predicted, mu_maps)
         loss.backward()
         return loss.item()
 
     def _eval_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
-        outpus = self.model(inputs)
-        loss = torch.nn.functional.loss.l1(outpus, mu_maps)
+        mu_maps_predicted = self.model(recons)
+        loss = torch.nn.functional.loss.l1(mu_maps_predicted, mu_maps)
         return loss.item()
 
 
@@ -144,6 +145,11 @@ if __name__ == "__main__":
         default="l1",
         help="define the loss function used for training, e.g. 0.75*l1+0.25*gdl",
     )
+    parser.add_argument(
+        "--decay_lr",
+        action="store_true",
+        help="decay the learning rate",
+    )
     parser.add_argument(
         "--lr", type=float, default=0.001, help="the initial learning rate for training"
     )
@@ -203,8 +209,6 @@ if __name__ == "__main__":
     torch.manual_seed(args.seed)
     np.random.seed(args.seed)
 
-    model = UNet(in_channels=1, features=args.features)
-    model = model.to(device)
 
     transform_normalization = None
     if args.input_norm == "mean":
@@ -214,41 +218,38 @@ if __name__ == "__main__":
     elif args.input_norm == "gaussian":
         transform_normalization = GaussianNormTransform()
 
-    data_loaders = {}
-    for split in ["train", "validation"]:
-        dataset = MuMapPatchDataset(
-            args.dataset_dir,
-            patches_per_image=args.number_of_patches,
-            patch_size=args.patch_size,
-            patch_offset=args.patch_offset,
-            shuffle=not args.no_shuffle,
-            split_name=split,
-            transform_normalization=transform_normalization,
-            logger=logger,
-        )
-        data_loader = torch.utils.data.DataLoader(
-            dataset=dataset,
-            batch_size=args.batch_size,
-            shuffle=True,
-            pin_memory=True,
-            num_workers=1,
+    dataset = MuMapPatchDataset(
+        args.dataset_dir,
+        patches_per_image=args.number_of_patches,
+        patch_size=args.patch_size,
+        patch_offset=args.patch_offset,
+        shuffle=not args.no_shuffle,
+        transform_normalization=transform_normalization,
+        logger=logger,
+    )
+
+    model = UNet(in_channels=1, features=args.features)
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.5, 0.999))
+    lr_scheduler = (
+        torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=args.lr_decay_factor, gamma=args.lr_decay_factor
         )
-        data_loaders[split] = data_loader
+        if args.decay_lr
+        else None
+    )
+    params = TrainingParams(name="Model", model=model, optimizer=optimizer, lr_scheduler=lr_scheduler)
 
     criterion = WeightedLoss.from_str(args.loss_func)
-    logger.debug(f"Criterion: {criterion}")
 
     training = Training(
-        model=model,
-        data_loaders=data_loaders,
         epochs=args.epochs,
+        dataset=dataset,
+        batch_size=args.batch_size,
         device=device,
-        loss_func=criterion,
-        lr=args.lr,
-        lr_decay_factor=args.lr_decay_factor,
-        lr_decay_epoch=args.lr_decay_epoch,
         snapshot_dir=args.snapshot_dir,
         snapshot_epoch=args.snapshot_epoch,
+        params=params,
+        loss_func=criterion,
         logger=logger,
     )
     training.run()

mu_map/training/lib.py

 from dataclasses import dataclass
 import os
-from typing import Dict
+from typing import Dict, Optional
+import sys
 
 import torch
 from torch import Tensor
 
 from mu_map.dataset.default import MuMapDataset
 
+
 @dataclass
 class TrainingParams:
     name: str
@@ -14,8 +16,8 @@ class TrainingParams:
     optimizer: torch.optim.Optimizer
     lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler]
 
-class AbstractTraining:
 
+class AbstractTraining:
     def __init__(
         self,
         epochs: int,
@@ -24,7 +26,7 @@ class AbstractTraining:
         device: torch.device,
         snapshot_dir: str,
         snapshot_epoch: int,
-        logger=None,
+        logger,  # TODO make optional?
     ):
         self.epochs = epochs
         self.batch_size = batch_size
@@ -37,7 +39,21 @@ class AbstractTraining:
         self.logger = logger
 
         self.training_params = []
-
+        self.data_loaders = dict(
+            [
+                (
+                    split_name,
+                    torch.utils.data.DataLoader(
+                        dataset.split_copy(split_name),
+                        batch_size=self.batch_size,
+                        shuffle=True,
+                        pin_memory=True,
+                        num_workers=1,
+                    ),
+                )
+                for split_name in ["train", "validation"]
+            ]
+        )
 
     def run(self) -> float:
         loss_val_min = sys.maxsize
@@ -88,7 +104,7 @@ class AbstractTraining:
             for param in self.training_params:
                 param.optimizer.zero_grad()
 
-            loss = loss + self._train_batch(self, inputs, targets)
+            loss = loss + self._train_batch(inputs, targets)
 
             for param in self.training_params:
                 param.optimizer.step()
@@ -110,12 +126,14 @@ class AbstractTraining:
             inputs = inputs.to(self.device)
             targets = targets.to(self.device)
 
-            loss = loss + self._eval_batch(self, inputs, targets)
+            loss = loss + self._eval_batch(inputs, targets)
         return loss / len(data_loader)
 
     def store_snapshot(self, prefix: str):
         for param in self.training_params:
-            snapshot_file = os.path.join(self.snapshot_dir, f"{prefix}_{param.name}.pth")
+            snapshot_file = os.path.join(
+                self.snapshot_dir, f"{prefix}_{param.name}.pth"
+            )
             self.logger.debug(f"Store snapshot at {snapshot_file}")
             torch.save(param.model.state_dict(), snapshot_file)
 
@@ -124,4 +142,3 @@ class AbstractTraining:
 
     def _eval_batch(self, inputs: torch.Tensor, targets: torch.Tensor) -> float:
         return 0
-