[common] Move reusable files to new common package

src/mednet/data/split.py → src/medbase/data/split.py

@@ -12,7 +12,7 @@ import typing
 
 import torch
 
-from .typing import DatabaseSplit, RawDataLoader
+from medbase.data.typing import DatabaseSplit, RawDataLoader
 
 logger = logging.getLogger(__name__)
 

src/mednet/utils/resources.py → src/medbase/utils/resources.py

@@ -18,7 +18,7 @@ import warnings
 import numpy
 import psutil
 
-from ..engine.device import SupportedPytorchDevice
+from medbase.engine.device import SupportedPytorchDevice
 
 logger = logging.getLogger(__name__)
 

src/medbase/utils/summary.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+# Adapted from https://github.com/pytorch/pytorch/issues/2001#issuecomment-405675488
+
+from functools import reduce
+
+import torch
+
+from torch.nn.modules.module import _addindent
+
+
+# ignore this space!
+def _repr(model: torch.nn.Module) -> tuple[str, int]:
+    # We treat the extra repr like the sub-module, one item per line
+    extra_lines = []
+    extra_repr = model.extra_repr()
+    # empty string will be split into list ['']
+    if extra_repr:
+        extra_lines = extra_repr.split("\n")
+    child_lines = []
+    total_params = 0
+    for key, module in model._modules.items():
+        mod_str, num_params = _repr(module)
+        mod_str = _addindent(mod_str, 2)
+        child_lines.append("(" + key + "): " + mod_str)
+        total_params += num_params
+    lines = extra_lines + child_lines
+
+    for _, p in model._parameters.items():
+        if hasattr(p, "dtype"):
+            total_params += reduce(lambda x, y: x * y, p.shape)
+
+    main_str = model._get_name() + "("
+    if lines:
+        # simple one-liner info, which most builtin Modules will use
+        if len(extra_lines) == 1 and not child_lines:
+            main_str += extra_lines[0]
+        else:
+            main_str += "\n  " + "\n  ".join(lines) + "\n"
+
+    main_str += ")"
+    main_str += f", {total_params:,} params"
+    return main_str, total_params
+
+
+def summary(model: torch.nn.Module) -> tuple[str, int]:
+    """Count the number of parameters in each model layer.
+
+    Parameters
+    ----------
+    model
+        Model to summarize.
+
+    Returns
+    -------
+    tuple[int, str]
+        A tuple containing a multiline string representation of the network and the number of parameters.
+    """
+    return _repr(model)

src/mednet/config/data/hivtb/datamodule.py

@@ -12,11 +12,12 @@ import pathlib
 import PIL.Image
 from torchvision.transforms.functional import to_tensor
 
-from ....data.datamodule import CachingDataModule
-from ....data.image_utils import remove_black_borders
-from ....data.split import make_split
-from ....data.typing import RawDataLoader as _BaseRawDataLoader
-from ....data.typing import Sample
+from medbase.data.datamodule import CachingDataModule
+from medbase.data.image_utils import remove_black_borders
+from medbase.data.split import make_split
+from medbase.data.typing import RawDataLoader as _BaseRawDataLoader
+from medbase.data.typing import Sample
+
 from ....utils.rc import load_rc
 
 CONFIGURATION_KEY_DATADIR = "datadir." + (__name__.rsplit(".", 2)[-2])

src/mednet/config/data/indian/datamodule.py

@@ -8,9 +8,10 @@ Database reference: [INDIAN-2013]_
 
 import pathlib
 
+from medbase.data.datamodule import CachingDataModule
+from medbase.data.split import make_split
+
 from ....config.data.shenzhen.datamodule import RawDataLoader
-from ....data.datamodule import CachingDataModule
-from ....data.split import make_split
 
 CONFIGURATION_KEY_DATADIR = "datadir." + (__name__.rsplit(".", 2)[-2])
 """Key to search for in the configuration file for the root directory of this

src/mednet/config/data/montgomery/datamodule.py

@@ -12,11 +12,12 @@ import pathlib
 import PIL.Image
 from torchvision.transforms.functional import to_tensor
 
-from ....data.datamodule import CachingDataModule
-from ....data.image_utils import remove_black_borders
-from ....data.split import make_split
-from ....data.typing import RawDataLoader as _BaseRawDataLoader
-from ....data.typing import Sample
+from medbase.data.datamodule import CachingDataModule
+from medbase.data.image_utils import remove_black_borders
+from medbase.data.split import make_split
+from medbase.data.typing import Sample
+from medbase.data.typing import RawDataLoader as _BaseRawDataLoader
+
 from ....utils.rc import load_rc
 
 CONFIGURATION_KEY_DATADIR = "datadir." + (__name__.rsplit(".", 2)[-2])