fix(train): ch folder struct, catch module errors

Author: okunator

cellseg_models_pytorch/training/__init__.py

@@ -1,13 +0,0 @@
-from .callbacks import METRIC_LOOKUP, Accuracy, MeanIoU
-from .lightning_experiment import SegmentationExperiment
-from .train_metrics import accuracy, confusion_mat, iou
-
-__all__ = [
-    "SegmentationExperiment",
-    "confusion_mat",
-    "accuracy",
-    "iou",
-    "Accuracy",
-    "MeanIoU",
-    "METRIC_LOOKUP",
-]

cellseg_models_pytorch/training/callbacks/metric_callbacks.py

@@ -1,21 +1,25 @@
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import torch
 
-from ..train_metrics import accuracy, iou
+from ..functional.train_metrics import accuracy, iou
 
 try:
     from torchmetrics import Metric
-except ImportError:
-    raise ImportError(
-        "`torchmetrics` package needed for metrics. `pip install torchmetrics`"
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "`torchmetrics` package is required when using metric callbacks. "
+        "Install with `pip install torchmetrics`"
     )
 
 
 __all__ = ["Accuracy", "MeanIoU"]
 
 
 class Accuracy(Metric):
+    higher_is_better: Optional[bool] = True
+    full_state_update: bool = False
+
     def __init__(
         self,
         compute_on_step: bool = True,
@@ -55,7 +59,7 @@ def update(
         self,
         pred: torch.Tensor,
         target: torch.Tensor,
-        activation: str = "sofmax",
+        activation: str = "softmax",
     ) -> None:
         """Update the batch accuracy list with one batch accuracy value.
 
@@ -83,6 +87,9 @@ def compute(self) -> torch.Tensor:
 
 
 class MeanIoU(Metric):
+    higher_is_better: Optional[bool] = True
+    full_state_update: bool = False
+
     def __init__(
         self,
         compute_on_step: bool = True,
@@ -119,7 +126,7 @@ def update(
         self,
         pred: torch.Tensor,
         target: torch.Tensor,
-        activation: str = "sofmax",
+        activation: str = "softmax",
     ) -> None:
         """Update the batch IoU list with one batch IoU matrix.
 

cellseg_models_pytorch/training/functional/__init__.py

@@ -0,0 +1,3 @@
+from .train_metrics import accuracy, confusion_mat, iou
+
+__all__ = ["confusion_mat", "accuracy", "iou"]

cellseg_models_pytorch/training/train_metrics.py renamed to cellseg_models_pytorch/training/functional/train_metrics.py

@@ -0,0 +1,3 @@
+from .lightning_experiment import SegmentationExperiment
+
+__all__ = ["SegmentationExperiment"]

cellseg_models_pytorch/training/lit/__init__.py

@@ -1,13 +1,20 @@
 from copy import deepcopy
 from typing import Any, Dict, List
 
-import pytorch_lightning as pl
 import torch
 import torch.nn as nn
 
-from ..losses import JOINT_SEG_LOSSES, SEG_LOSS_LOOKUP, JointLoss, Loss, MultiTaskLoss
-from ..optimizers import OPTIM_LOOKUP, SCHED_LOOKUP, adjust_optim_params
-from .callbacks import METRIC_LOOKUP
+try:
+    import pytorch_lightning as pl
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "To use the `SegmentationExperiment`, pytorch-lightning is required. "
+        "Install with `pip install pytorch-lightning`"
+    )
+
+from ...losses import JOINT_SEG_LOSSES, SEG_LOSS_LOOKUP, JointLoss, Loss, MultiTaskLoss
+from ...optimizers import OPTIM_LOOKUP, SCHED_LOOKUP, adjust_optim_params
+from ..callbacks import METRIC_LOOKUP
 
 
 class SegmentationExperiment(pl.LightningModule):

cellseg_models_pytorch/training/lightning_experiment.py renamed to cellseg_models_pytorch/training/lit/lightning_experiment.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from cellseg_models_pytorch.training.train_metrics import accuracy, iou
+from cellseg_models_pytorch.training.functional.train_metrics import accuracy, iou
 
 
 @pytest.mark.parametrize("metric", [accuracy, iou])

cellseg_models_pytorch/training/tests/test_metrics.py

@@ -6,7 +6,7 @@
 from cellseg_models_pytorch.datamodules.custom_datamodule import CustomDataModule
 from cellseg_models_pytorch.datasets import SegmentationFolderDataset
 from cellseg_models_pytorch.models import cellpose_plus
-from cellseg_models_pytorch.training import SegmentationExperiment
+from cellseg_models_pytorch.training.lit import SegmentationExperiment
 
 
 # @pytest.mark.parametrize