konfai 1.0.9__tar.gz → 1.1.0__tar.gz

{konfai-1.0.9 → konfai-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.0.9
+Version: 1.1.0
 Summary: Modular and configurable Deep Learning framework with YAML and PyTorch
 Author-email: Valentin Boussot <boussot.v@gmail.com>
 License-Expression: Apache-2.0

{konfai-1.0.9 → konfai-1.1.0}/konfai/data/data_manager.py

@@ -257,14 +257,14 @@ class Data(ABC):
                         subset : Union[Subset, dict[str, Subset]],
                         num_workers : int,
                         batch_size : int,
-                        train_size: Union[float, str, list[int], list[str]] = 1,
+                        validation: Union[float, str, list[int], list[str]] = 1,
                         inlineAugmentations: bool = False,
                         dataAugmentationsList: dict[str, DataAugmentationsList]= {}) -> None:
         self.dataset_filenames = dataset_filenames
         self.subset = subset
         self.groups_src = groups_src
         self.patch = patch
-        self.train_size = train_size
+        self.validation = validation
         self.dataAugmentationsList = dataAugmentationsList
         self.batch_size = batch_size
         self.dataSet_args = dict(groups_src=self.groups_src, inlineAugmentations=inlineAugmentations, dataAugmentationsList = list(self.dataAugmentationsList.values()), use_cache = use_cache, buffer_size=batch_size+1, patch_size=self.patch.patch_size if self.patch is not None else None, overlap=self.patch.overlap if self.patch is not None else None)
@@ -372,17 +372,17 @@ class Data(ABC):
     
         train_map = map
         validate_map = []
-        if isinstance(self.train_size, float):
-            if self.train_size < 1.0 and int(math.floor(len(map)*(1-self.train_size))) > 0:
-                train_map, validate_map = map[:int(math.floor(len(map)*self.train_size))], map[int(math.floor(len(map)*self.train_size)):]
-        elif isinstance(self.train_size, str):
-            if ":" in self.train_size:
+        if isinstance(self.validation, float):
+            if self.validation < 1.0 and int(math.floor(len(map)*(1-self.validation))) > 0:
+                train_map, validate_map = map[:int(math.floor(len(map)*self.validation))], map[int(math.floor(len(map)*self.validation)):]
+        elif isinstance(self.validation, str):
+            if ":" in self.validation:
                 index = list(range(int(self.subset.split(":")[0]), int(self.subset.split(":")[1])))
                 train_map = [m for m in map if m[0] not in index]
                 validate_map = [m for m in map if m[0] in index]
-            elif os.path.exists(self.train_size):
+            elif os.path.exists(self.validation):
                 validation_names = []
-                with open(self.train_size, "r") as f:
+                with open(self.validation, "r") as f:
                     for name in f:
                         validation_names.append(name.strip())
                 index = [i for i, n in enumerate(subset_names) if n in validation_names]
@@ -390,13 +390,13 @@ class Data(ABC):
                 validate_map = [m for m in map if m[0] in index]
             else:
                 validate_map = train_map
-        elif isinstance(self.train_size, list):
-            if len(self.train_size) > 0:
-                if isinstance(self.train_size[0], int):
-                    train_map = [m for m in map if m[0] not in self.train_size]
-                    validate_map = [m for m in map if m[0] in self.train_size]
-                elif isinstance(self.train_size[0], str):
-                    index = [i for i, n in enumerate(subset_names) if n in self.train_size]
+        elif isinstance(self.validation, list):
+            if len(self.validation) > 0:
+                if isinstance(self.validation[0], int):
+                    train_map = [m for m in map if m[0] not in self.validation]
+                    validate_map = [m for m in map if m[0] in self.validation]
+                elif isinstance(self.validation[0], str):
+                    index = [i for i, n in enumerate(subset_names) if n in self.validation]
                     train_map = [m for m in map if m[0] not in index]
                     validate_map = [m for m in map if m[0] in index]
 
@@ -436,8 +436,8 @@ class DataTrain(Data):
                         subset : Union[TrainSubset, dict[str, TrainSubset]] = TrainSubset(),
                         num_workers : int = 4,
                         batch_size : int = 1,
-                        train_size : Union[float, str, list[int], list[str]] = 0.8) -> None:
-        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, train_size, inlineAugmentations, augmentations if augmentations else {})
+                        validation : Union[float, str, list[int], list[str]] = 0.8) -> None:
+        super().__init__(dataset_filenames, groups_src, patch, use_cache, subset, num_workers, batch_size, validation, inlineAugmentations, augmentations if augmentations else {})
 
 class DataPrediction(Data):
 
@@ -463,4 +463,4 @@ class DataMetric(Data):
                         validation: Union[str, None] = None,
                         num_workers : int = 4) -> None:
 
-        super().__init__(dataset_filenames=dataset_filenames, groups_src=groups_src, patch=None, use_cache=False, subset=subset, num_workers=num_workers, batch_size=1, train_size=1 if validation is None else validation)
+        super().__init__(dataset_filenames=dataset_filenames, groups_src=groups_src, patch=None, use_cache=False, subset=subset, num_workers=num_workers, batch_size=1, validation=1 if validation is None else validation)

{konfai-1.0.9 → konfai-1.1.0}/konfai/data/patching.py

@@ -188,7 +188,7 @@ class Patch(ABC):
             pad_bottom = 0
             pad_top = 0
             if self._patch_slices[a][index][0].start-bottom < 0:
-                pad_bottom = bottom-s.start
+                pad_bottom = bottom-self._patch_slices[a][index][0].start
             if self._patch_slices[a][index][0].stop+top > data.shape[len(slices_pre)]:
                 pad_top = self._patch_slices[a][index][0].stop+top-data.shape[len(slices_pre)]
             data_sliced = F.pad(data_sliced, [0 for _ in range((len(slices)-1)*2)]+[pad_bottom, pad_top], 'reflect')

{konfai-1.0.9 → konfai-1.1.0}/konfai/metric/measure.py

@@ -147,7 +147,7 @@ class Dice(Criterion):
     def forward(self, output: torch.Tensor, *targets : list[torch.Tensor]) -> torch.Tensor:
         target = targets[0]
         if output.shape[1] == 1:
-            output = F.one_hot(output.type(torch.int64), num_classes=torch.max(output).item()+1).permute(0, len(target.shape), *[i+1 for i in range(len(target.shape)-1)]).float()
+            output = F.one_hot(output.type(torch.int64), num_classes=int(torch.max(output).item()+1)).permute(0, len(target.shape), *[i+1 for i in range(len(target.shape)-1)]).float()
         target = F.one_hot(target.type(torch.int64), num_classes=output.shape[1]).permute(0, len(target.shape), *[i+1 for i in range(len(target.shape)-1)]).float().squeeze(2)
         return 1-torch.mean(self.dice_per_channel(output, target))
 

{konfai-1.0.9 → konfai-1.1.0}/konfai/trainer.py

@@ -1,5 +1,4 @@
 import torch
-import torch.optim.adamw
 from torch.utils.data import DataLoader
 import tqdm
 import numpy as np
@@ -19,9 +18,72 @@ from konfai.utils.utils import State, DataLog, DistributedObject, description
 from konfai.network.network import Network, ModelLoader, NetState, CPU_Model
 
 
+class EarlyStoppingBase:
+
+    def __init__(self):
+        pass
+
+    def isStopped(self) -> bool:
+        return False
+
+    def getScore(self, values: dict[str, float]):
+        return sum([i for i in values.values()])
+        
+    def __call__(self, current_score: float) -> bool:
+        return False
+
+class EarlyStopping(EarlyStoppingBase):
+
+    @config("EarlyStopping")
+    def __init__(self, monitor: Union[list[str], None] = [], patience=10, min_delta=0.0, mode="min"):
+        super().__init__()
+        self.monitor = [] if monitor is None else monitor
+        self.patience = patience
+        self.min_delta = min_delta
+        self.mode = mode
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+
+    def isStopped(self) -> bool:
+        return self.early_stop
+
+    def getScore(self, values: dict[str, float]):
+        if len(self.monitor) == 0:
+            return super().getScore(values)
+        for v in self.monitor:
+            if v not in values.keys():
+                raise ValueError(
+                    "[EarlyStopping] Metric '{}' specified in `monitor` not found in logged values. "
+                    "Available keys: {}. Please check your configuration.".format(v, list(values.keys())))
+        return sum([i for v, i in values.items() if v in self.monitor])
+        
+    def __call__(self, current_score: float) -> bool:
+        if self.best_score is None:
+            self.best_score = current_score
+            return False
+
+        if self.mode == "min":
+            improvement = self.best_score - current_score
+        elif self.mode == "max":
+            improvement = current_score - self.best_score
+        else:
+            raise ValueError("Mode must be 'min' or 'max'.")
+
+        if improvement > self.min_delta:
+            self.best_score = current_score
+            self.counter = 0
+        else:
+            self.counter += 1
+
+        if self.counter >= self.patience:
+            self.early_stop = True
+
+        return self.early_stop
+
 class _Trainer():
 
-    def __init__(self, world_size: int, global_rank: int, local_rank: int, size: int, train_name: str, data_log: Union[list[str], None] , save_checkpoint_mode: str, epochs: int, epoch: int, autocast: bool, it_validation: Union[int, None], it: int, model: Union[DDP, CPU_Model], modelEMA: AveragedModel, dataloader_training: DataLoader, dataloader_validation: Union[DataLoader, None] = None) -> None:
+    def __init__(self, world_size: int, global_rank: int, local_rank: int, size: int, train_name: str, early_stopping: EarlyStopping, data_log: Union[list[str], None] , save_checkpoint_mode: str, epochs: int, epoch: int, autocast: bool, it_validation: Union[int, None], it: int, model: Union[DDP, CPU_Model], modelEMA: AveragedModel, dataloader_training: DataLoader, dataloader_validation: Union[DataLoader, None] = None) -> None:
         self.world_size = world_size        
         self.global_rank = global_rank
         self.local_rank = local_rank
@@ -36,7 +98,8 @@ class _Trainer():
         self.dataloader_validation = dataloader_validation
         self.autocast = autocast
         self.modelEMA = modelEMA
-        
+        self.early_stopping = EarlyStoppingBase() if early_stopping is None else early_stopping 
+
         self.it_validation = it_validation
         if self.it_validation is None:
             self.it_validation = len(dataloader_training)
@@ -59,6 +122,8 @@ class _Trainer():
             for self.epoch in epoch_tqdm:
                 self.dataloader_training.dataset.load()
                 self.train()
+                if self.early_stopping.isStopped():
+                    break
                 self.dataloader_training.dataset.resetAugmentation()
                 
     def getInput(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int, str, bool]]) -> dict[tuple[str, bool], torch.Tensor]:
@@ -88,7 +153,11 @@ class _Trainer():
                         if self.dataloader_validation is not None:
                             loss = self._validate()
                         self.model.module.update_lr()
-                        self.checkpoint_save(loss)
+                        score = self.early_stopping.getScore(loss)
+                        self.checkpoint_save(score)
+                        if self.early_stopping(score):
+                            break
+                        
 
                 batch_iter.set_description(desc()) 
             
@@ -120,41 +189,53 @@ class _Trainer():
         return self._validation_log(data_dict)
 
     def checkpoint_save(self, loss: float) -> None:
-        if self.global_rank == 0:
-            path = CHECKPOINTS_DIRECTORY()+self.train_name+"/"
-            last_loss = None
-            if os.path.exists(path) and os.listdir(path):
-                name = sorted(os.listdir(path))[-1]
-                state_dict = torch.load(path+name, weights_only=False)
-                last_loss = state_dict["loss"]
-                if self.save_checkpoint_mode == "BEST":
-                    if last_loss >= loss:
-                        os.remove(path+name)
-
-            if self.save_checkpoint_mode != "BEST" or (last_loss is None or last_loss >= loss):
-                name = DATE()+".pt"
-                if not os.path.exists(path):
-                    os.makedirs(path)
-                
-                save_dict = {
-                    "epoch": self.epoch,
-                    "it": self.it,
-                    "loss": loss,
-                    "Model": self.model.module.state_dict()}
+        if self.global_rank != 0:
+            return
 
-                if self.modelEMA is not None:
-                    save_dict.update({"Model_EMA" : self.modelEMA.module.state_dict()})
+        path = CHECKPOINTS_DIRECTORY()+self.train_name+"/"
+        os.makedirs(path, exist_ok=True)
+
+        name = DATE() + ".pt"
+        save_path = os.path.join(path, name)
 
-                save_dict.update({'{}_optimizer_state_dict'.format(name): network.optimizer.state_dict() for name, network in self.model.module.getNetworks().items() if network.optimizer is not None})
-                torch.save(save_dict, path+name)
+        save_dict = {
+        "epoch": self.epoch,
+        "it": self.it,
+        "loss": loss,
+        "Model": self.model.module.state_dict()
+        }
+
+        if self.modelEMA is not None:
+            save_dict["Model_EMA"] = self.modelEMA.module.state_dict()
+        
+        save_dict.update({'{}_optimizer_state_dict'.format(name): network.optimizer.state_dict() for name, network in self.model.module.getNetworks().items() if network.optimizer is not None})
+        torch.save(save_dict, save_path)
+
+        if self.save_checkpoint_mode == "BEST":
+            all_checkpoints = sorted([
+                os.path.join(path, f)
+                for f in os.listdir(path) if f.endswith(".pt")
+            ])
+            best_ckpt = None
+            best_loss = float('inf')
+
+            for f in all_checkpoints:
+                d = torch.load(f, weights_only=False)
+                if d.get("loss", float("inf")) < best_loss:
+                    best_loss = d["loss"]
+                    best_ckpt = f
+
+            for f in all_checkpoints:
+                if f != best_ckpt and f != save_path:
+                    os.remove(f)
 
     @torch.no_grad()
-    def _log(self, type_log: str, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> float:
+    def _log(self, type_log: str, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> dict[str, float]:
         models: dict[str, Network] = {"" : self.model.module}
         if self.modelEMA is not None:
             models["_EMA"] = self.modelEMA.module
         
-        measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank*self.size+self.size-1, models, self.it_validation if type_log == "Trainning" else len(self.dataloader_validation))
+        measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank*self.size+self.size-1, models, self.it_validation if type_log == "Training" else len(self.dataloader_validation))
         
         if self.global_rank == 0:
             images_log = []
@@ -178,27 +259,29 @@ class _Trainer():
                         for name, layer, _ in model.get_layers([v.to(0) for k, v in self.getInput(data_dict).items() if k[1]], images_log):
                             self.data_log[name][0](self.tb, "{}/{}{}".format(type_log, name, label), layer[:self.data_log[name][1]].detach().cpu().numpy(), self.it)
                         
-            if type_log == "Trainning":
+            if type_log == "Training":
                 for name, network in self.model.module.getNetworks().items():
                     if network.optimizer is not None:
                         self.tb.add_scalar("{}/{}/Learning Rate".format(type_log, name), network.optimizer.param_groups[0]['lr'], self.it)
         
         if self.global_rank == 0:
-            loss = []
+            loss = {}
             for name, network in self.model.module.getNetworks().items():
                 if network.measure is not None:
-                    loss.append(sum([v[1] for v in measures["{}".format(name)][0].values()]))
-            return np.mean(loss)
+                    loss.update({k : v[1] for k, v in measures["{}{}".format(name, label)][0].items()})
+                    loss.update({k : v[1] for k, v in measures["{}{}".format(name, label)][1].items()})
+            return loss
         return None
     
     @torch.no_grad()
-    def _train_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> float:
-        return self._log("Trainning", data_dict)
+    def _train_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> dict[str, float]:
+        return self._log("Training", data_dict)
 
     @torch.no_grad()
-    def _validation_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> float:
+    def _validation_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]) -> dict[str, float]:
         return self._log("Validation", data_dict)
 
+    
 class Trainer(DistributedObject):
 
     @config("Trainer")
@@ -214,6 +297,7 @@ class Trainer(DistributedObject):
                     gpu_checkpoints: Union[list[str], None] = None,
                     ema_decay : float = 0,
                     data_log: Union[list[str], None] = None,
+                    early_stopping: Union[EarlyStopping, None] = None,
                     save_checkpoint_mode: str= "BEST") -> None:
         if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
@@ -223,6 +307,7 @@ class Trainer(DistributedObject):
         self.autocast = autocast
         self.epochs = epochs
         self.epoch = 0
+        self.early_stopping = early_stopping
         self.it = 0
         self.it_validation = it_validation
         self.model = model.getModel(train=True)
@@ -326,5 +411,5 @@ class Trainer(DistributedObject):
         model = DDP(model, static_graph=True) if torch.cuda.is_available() else CPU_Model(model)
         if self.modelEMA is not None:
             self.modelEMA.module = Network.to(self.modelEMA.module, local_rank)
-        with _Trainer(world_size, global_rank, local_rank, self.size, self.name, self.data_log, self.save_checkpoint_mode, self.epochs, self.epoch, self.autocast, self.it_validation, self.it, model, self.modelEMA, *dataloaders) as t:
+        with _Trainer(world_size, global_rank, local_rank, self.size, self.name, self.early_stopping, self.data_log, self.save_checkpoint_mode, self.epochs, self.epoch, self.autocast, self.it_validation, self.it, model, self.modelEMA, *dataloaders) as t:
             t.run()

{konfai-1.0.9 → konfai-1.1.0}/konfai/utils/config.py

@@ -236,7 +236,7 @@ def config(key : Union[str, None] = None):
                                     try:
                                         kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
                                     except Exception as e:
-                                        raise ValueError("[Config] Failed to instantiate {} with type {}".format(param.name, annotation.__name__))
+                                        raise ValueError("[Config] Failed to instantiate {} with type {}, error {} ".format(param.name, annotation.__name__, e))
                                     
                                     if os.environ['KONFAI_CONFIG_VARIABLE'] == "True":
                                         os.environ['KONFAI_CONFIG_VARIABLE'] = "False"

{konfai-1.0.9 → konfai-1.1.0}/konfai.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.0.9
+Version: 1.1.0
 Summary: Modular and configurable Deep Learning framework with YAML and PyTorch
 Author-email: Valentin Boussot <boussot.v@gmail.com>
 License-Expression: Apache-2.0

{konfai-1.0.9 → konfai-1.1.0}/konfai.egg-info/SOURCES.txt

@@ -39,4 +39,5 @@ konfai/utils/__init__.py
 konfai/utils/config.py
 konfai/utils/dataset.py
 konfai/utils/registration.py
-konfai/utils/utils.py
+konfai/utils/utils.py
+tests/test_config.py

{konfai-1.0.9 → konfai-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "konfai"
-version = "1.0.9"
+version = "1.1.0"
 description = "Modular and configurable Deep Learning framework with YAML and PyTorch"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -41,4 +41,7 @@ konfai-cluster = "konfai.main:cluster"
 vtk = ["vtk"]
 lpips = ["lpips"]
 cluster = ["submitit"]
-plot = ["matplotlib"]
+plot = ["matplotlib"]
+
+[tool.poetry.dev-dependencies]
+pytest = "^8.0"

konfai-1.1.0/tests/test_config.py

@@ -0,0 +1,18 @@
+import pytest
+from konfai.utils.config import config
+
+class Dummy:
+
+    @config("Test")
+    def __init__(self, a: int = 1, b: str = "ok"):
+        self.a = a
+        self.b = b
+
+def test_config_instantiation(monkeypatch):
+    import os
+    os.environ['KONFAI_CONFIG_FILE'] = "./tests/dummy_data/dummy_config.yml"
+    os.environ['KONFAI_CONFIG_PATH'] = "Test"
+
+    dummy = Dummy()
+    assert dummy.a == 42
+    assert dummy.b == "hello"