konfai 1.0.8__py3-none-any.whl → 1.1.0__py3-none-any.whl

konfai/network/network.py

@@ -13,11 +13,11 @@ from torch.utils.checkpoint import checkpoint
 from typing import Union
 from enum import Enum
 
-from konfai import DEEP_LEARNING_API_ROOT
+from konfai import KONFAI_ROOT
 from konfai.metric.schedulers import Scheduler
 from konfai.utils.config import config
 from konfai.utils.utils import State, _getModule, getDevice, getGPUMemory
-from konfai.data.HDF5 import Accumulator, ModelPatch
+from konfai.data.patching import Accumulator, ModelPatch
 
 class NetState(Enum):
     TRAIN = 0,
@@ -40,7 +40,7 @@ class OptimizerLoader():
     
     def getOptimizer(self, key: str, parameter: Iterator[torch.nn.parameter.Parameter]) -> torch.optim.Optimizer:
         torch.optim.AdamW
-        return config("{}.Model.{}.Optimizer".format(DEEP_LEARNING_API_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
+        return config("{}.Model.{}.Optimizer".format(KONFAI_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
         
 class SchedulerStep():
     
@@ -98,8 +98,8 @@ class CriterionsLoader():
         for module_classpath, criterionsAttr in self.criterionsLoader.items():
             module, name = _getModule(module_classpath, "metric.measure")
             criterionsAttr.isTorchCriterion = module.startswith("torch")
-            criterionsAttr.sheduler = criterionsAttr.l.getShedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(DEEP_LEARNING_API_ROOT(), model_classname, output_group, target_group, module_classpath))
-            criterions[config("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(DEEP_LEARNING_API_ROOT(), model_classname, output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
+            criterionsAttr.sheduler = criterionsAttr.l.getShedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))
+            criterions[config("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
         return criterions
 
 class TargetCriterionsLoader():
@@ -753,14 +753,14 @@ class Network(ModuleArgsDict, ABC):
         output_layer_accumulator : dict[str, Accumulator] = {}
         output_layer_patch_indexed : dict[str, Patch_Indexed] = {}
         it = 0
-        debug = "DL_API_DEBUG" in os.environ
+        debug = "KONFAI_DEBUG" in os.environ
         for (nameTmp, output_layer) in self.named_forward(*inputs):
             name = nameTmp.replace(";accu;", "")
             if debug:
-                if "DL_API_DEBUG_LAST_LAYER" in os.environ:
-                    os.environ["DL_API_DEBUG_LAST_LAYER"] = "{}|{}:{}:{}".format(os.environ["DL_API_DEBUG_LAST_LAYER"], name, getGPUMemory(output_layer.device), str(output_layer.device).replace("cuda:", ""))
+                if "KONFAI_DEBUG_LAST_LAYER" in os.environ:
+                    os.environ["KONFAI_DEBUG_LAST_LAYER"] = "{}|{}:{}:{}".format(os.environ["KONFAI_DEBUG_LAST_LAYER"], name, getGPUMemory(output_layer.device), str(output_layer.device).replace("cuda:", ""))
                 else:
-                    os.environ["DL_API_DEBUG_LAST_LAYER"] = "{}:{}:{}".format(name, getGPUMemory(output_layer.device), str(output_layer.device).replace("cuda:", ""))
+                    os.environ["KONFAI_DEBUG_LAST_LAYER"] = "{}:{}:{}".format(name, getGPUMemory(output_layer.device), str(output_layer.device).replace("cuda:", ""))
             it += 1
             if name in layers_name or nameTmp in layers_name:
                 if ";accu;" in nameTmp:
@@ -918,12 +918,12 @@ class Network(ModuleArgsDict, ABC):
 class ModelLoader():
 
     @config("Model")
-    def __init__(self, classpath : str = "default:segmentation.UNet") -> None:
-        self.module, self.name = _getModule(classpath.split(".")[-1] if len(classpath.split(".")) > 1 else classpath, ".".join(classpath.split(".")[:-1]) if len(classpath.split(".")) > 1 else "")
+    def __init__(self, classpath : str = "default:segmentation.UNet.UNet") -> None:
+        self.module, self.name = _getModule(classpath, "models")
         
     def getModel(self, train : bool = True, DL_args: Union[str, None] = None, DL_without=["optimizer", "schedulers", "nb_batch_per_step", "init_type", "init_gain"]) -> Network:
         if not DL_args:
-            DL_args="{}.Model".format(DEEP_LEARNING_API_ROOT())
+            DL_args="{}.Model".format(KONFAI_ROOT())
         model = partial(getattr(importlib.import_module(self.module), self.name), config = None, DL_args=DL_args)
         if not train: 
             model = partial(model, DL_without = DL_without)

konfai/predictor.py

@@ -6,12 +6,12 @@ import torch
 import tqdm
 import os
 
-from konfai import MODELS_DIRECTORY, PREDICTIONS_DIRECTORY, CONFIG_FILE, MODEL, DEEP_LEARNING_API_ROOT
+from konfai import MODELS_DIRECTORY, PREDICTIONS_DIRECTORY, CONFIG_FILE, MODEL, KONFAI_ROOT
 from konfai.utils.config import config
 from konfai.utils.utils import State, get_patch_slices_from_nb_patch_per_dim, NeedDevice, _getModule, DistributedObject, DataLog, description
 from konfai.utils.dataset import Dataset, Attribute
-from konfai.data.dataset import DataPrediction, DatasetIter
-from konfai.data.HDF5 import Accumulator, PathCombine
+from konfai.data.data_manager import DataPrediction, DatasetIter
+from konfai.data.patching import Accumulator, PathCombine
 from konfai.network.network import ModelLoader, Network, NetState, CPU_Model
 from konfai.data.transform import Transform, TransformLoader
 
@@ -52,13 +52,13 @@ class OutDataset(Dataset, NeedDevice, ABC):
             
             if _transform_type is not None:
                 for classpath, transform in _transform_type.items():
-                    transform = transform.getTransform(classpath, DL_args =  "{}.outsDataset.{}.OutDataset.{}".format(DEEP_LEARNING_API_ROOT(), name_layer, name))
+                    transform = transform.getTransform(classpath, DL_args =  "{}.outsDataset.{}.OutDataset.{}".format(KONFAI_ROOT(), name_layer, name))
                     transform.setDatasets(datasets)
                     transform_type.append(transform)
 
         if self._patchCombine is not None:
-            module, name = _getModule(self._patchCombine, "data.HDF5")
-            self.patchCombine = getattr(importlib.import_module(module), name)(config = None, DL_args =  "{}.outsDataset.{}.OutDataset".format(DEEP_LEARNING_API_ROOT(), name_layer))
+            module, name = _getModule(self._patchCombine, "data.patching")
+            self.patchCombine = getattr(importlib.import_module(module), name)(config = None, DL_args =  "{}.outsDataset.{}.OutDataset".format(KONFAI_ROOT(), name_layer))
     
     def setPatchConfig(self, patchSize: Union[list[int], None], overlap: Union[int, None], nb_data_augmentation: int) -> None:
         if patchSize is not None and overlap is not None:
@@ -94,7 +94,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
 class OutSameAsGroupDataset(OutDataset):
 
     @config("OutDataset")
-    def __init__(self, dataset_filename: str = "Dataset:h5", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, redution: str = "mean", inverse_transform: bool = True) -> None:
+    def __init__(self, dataset_filename: str = "./Dataset:mha", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, redution: str = "mean", inverse_transform: bool = True) -> None:
         super().__init__(dataset_filename, group, pre_transforms, post_transforms, final_transforms, patchCombine)
         self.group_src, self.group_dest = sameAsGroup.split(":")
         self.redution = redution
@@ -240,7 +240,7 @@ class _Predictor():
         self.modelComposite.module.setState(NetState.PREDICTION)
         desc = lambda : "Prediction : {}".format(description(self.modelComposite))
         self.dataloader_prediction.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=False, desc = desc(), total=len(self.dataloader_prediction), disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=False, desc = desc(), total=len(self.dataloader_prediction), disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             dist.barrier()
             for it, data_dict in batch_iter:
                 input = self.getInput(data_dict)
@@ -322,7 +322,7 @@ class Predictor(DistributedObject):
                     gpu_checkpoints: Union[list[str], None] = None,
                     outsDataset: Union[dict[str, OutDatasetLoader], None] = {"default:Default" : OutDatasetLoader()},
                     images_log: list[str] = []) -> None:
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] != "Done":
+        if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
         super().__init__(train_name)
         self.manual_seed = manual_seed
@@ -374,7 +374,7 @@ class Predictor(DistributedObject):
         for dataset_filename in self.datasets_filename:
             path = self.predict_path +dataset_filename
             if os.path.exists(path):
-                if os.environ["DL_API_OVERWRITE"] != "True":
+                if os.environ["KONFAI_OVERWRITE"] != "True":
                     accept = builtins.input("The prediction {} already exists ! Do you want to overwrite it (yes,no) : ".format(path))
                     if accept != "yes":
                         return

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
@@ -12,16 +11,79 @@ from torch.utils.tensorboard.writer import SummaryWriter
 from torch.optim.swa_utils import AveragedModel
 import torch.distributed as dist
 
-from konfai import MODELS_DIRECTORY, CHECKPOINTS_DIRECTORY, STATISTICS_DIRECTORY, SETUPS_DIRECTORY, CONFIG_FILE, MODEL, DATE, DL_API_STATE
-from konfai.data.dataset import DataTrain
+from konfai import MODELS_DIRECTORY, CHECKPOINTS_DIRECTORY, STATISTICS_DIRECTORY, SETUPS_DIRECTORY, CONFIG_FILE, MODEL, DATE, KONFAI_STATE
+from konfai.data.data_manager import DataTrain
 from konfai.utils.config import config
 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]:
@@ -72,7 +137,7 @@ class _Trainer():
             self.modelEMA.module.setState(NetState.TRAIN)
 
         desc = lambda : "Training : {}".format(description(self.model, self.modelEMA))
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             for _, data_dict in batch_iter:
                 with torch.amp.autocast('cuda', enabled=self.autocast):
                     input = self.getInput(data_dict)
@@ -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()) 
             
@@ -103,7 +172,7 @@ class _Trainer():
         desc = lambda : "Validation : {}".format(description(self.model, self.modelEMA))
         data_dict = None
         self.dataloader_validation.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             for _, data_dict in batch_iter:
                 input = self.getInput(data_dict)
                 self.model(input)
@@ -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,34 +259,36 @@ 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")
     def __init__(   self,
                     model : ModelLoader = ModelLoader(),
                     dataset : DataTrain = DataTrain(),
-                    train_name : str = "default:name",
+                    train_name : str = "default:TRAIN_01",
                     manual_seed : Union[int, None] = None,
                     epochs: int = 100,
                     it_validation : Union[int, None] = None,
@@ -214,8 +297,9 @@ 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["DEEP_LEANING_API_CONFIG_MODE"] != "Done":
+        if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
         super().__init__(train_name)
         self.manual_seed = manual_seed        
@@ -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)
@@ -292,9 +377,9 @@ class Trainer(DistributedObject):
         return (1-self.ema_decay) * averaged_model_parameter + self.ema_decay * model_parameter
     
     def setup(self, world_size: int):
-        state = State._member_map_[DL_API_STATE()]
-        if state != State.RESUME and os.path.exists(STATISTICS_DIRECTORY()+self.name+"/"):
-            if os.environ["DL_API_OVERWRITE"] != "True":
+        state = State._member_map_[KONFAI_STATE()]
+        if state != State.RESUME and os.path.exists(CHECKPOINTS_DIRECTORY()+self.name+"/"):
+            if os.environ["KONFAI_OVERWRITE"] != "True":
                 accept = input("The model {} already exists ! Do you want to overwrite it (yes,no) : ".format(self.name))
                 if accept != "yes":
                     return
@@ -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/utils/config.py

@@ -3,9 +3,8 @@ import ruamel.yaml
 import inspect 
 import collections
 from copy import deepcopy
-from typing import Union
+from typing import Union, Literal, get_origin, get_args
 import torch
-
 from konfai import CONFIG_FILE
 
 yaml = ruamel.yaml.YAML()
@@ -26,7 +25,7 @@ class Config():
         if not os.path.exists(self.filename):
             result = input("Create a new config file ? [no,yes,interactive] : ")
             if result in ["yes", "interactive"]:
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "interactive" if result == "interactive" else "default"
+                os.environ["KONFAI_CONFIG_MODE"] = "interactive" if result == "interactive" else "default"
             else:
                 exit(0)
             with open(self.filename, "w") as f:
@@ -69,7 +68,7 @@ class Config():
 
     def __exit__(self, type, value, traceback) -> None:
         self.yml.close()
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "remove":
+        if os.environ["KONFAI_CONFIG_MODE"] == "remove":
             if os.path.exists(CONFIG_FILE()):
                 os.remove(CONFIG_FILE())
             return
@@ -87,22 +86,22 @@ class Config():
         except:
             result = input("\nKeep a default configuration file ? (yes,no) : ")
             if result == "yes":
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "default"
+                os.environ["KONFAI_CONFIG_MODE"] = "default"
             else:
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "remove"
+                os.environ["KONFAI_CONFIG_MODE"] = "remove"
                 exit(0)
         return default.split(":")[1] if len(default.split(":")) > 1 else default
 
     @staticmethod
     def _getInputDefault(name : str, default : Union[str, None], isList : bool = False) -> Union[list[Union[str, None]], str, None]:
         if isinstance(default, str) and (default == "default" or (len(default.split(":")) > 1 and default.split(":")[0] == "default")):
-            if os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "interactive":
+            if os.environ["KONFAI_CONFIG_MODE"] == "interactive":
                 if isList:
                     list_tmp = []
                     key_tmp = "OK"
-                    while key_tmp != "!" and os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "interactive":
+                    while (key_tmp != "!" and key_tmp != " ") and os.environ["KONFAI_CONFIG_MODE"] == "interactive":
                         key_tmp = Config._getInput(name, default)
-                        if key_tmp != "!":
+                        if (key_tmp != "!" and key_tmp != " "):
                             if key_tmp == "":
                                 key_tmp = default.split(":")[1] if len(default.split(":")) > 1 else default
                             list_tmp.append(key_tmp)
@@ -134,6 +133,7 @@ class Config():
                 list_tmp = [] 
                 for key in value_config:
                     list_tmp.extend(Config._getInputDefault(name, key, isList=True))
+                
 
                 value = list_tmp
                 value_config = list_tmp
@@ -155,7 +155,7 @@ class Config():
                     dict_value[key] = value_tmp
                 value = dict_value
         if isinstance(self.config, str):
-            os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "True"
+            os.environ['KONFAI_CONFIG_VARIABLE'] = "True"
             return None
 
         self.config[name] = value_config if value_config is not None else "None"
@@ -169,17 +169,30 @@ def config(key : Union[str, None] = None):
             if "config" in kwargs:
                 filename =  kwargs["config"]
                 if filename == None:
-                    filename = os.environ['DEEP_LEARNING_API_CONFIG_FILE']
+                    filename = os.environ['KONFAI_CONFIG_FILE']
                 else:
-                    os.environ['DEEP_LEARNING_API_CONFIG_FILE'] = filename
+                    os.environ['KONFAI_CONFIG_FILE'] = filename
                 key_tmp =  kwargs["DL_args"]+("."+key if key is not None else "") if "DL_args" in kwargs else key
                 without =  kwargs["DL_without"] if "DL_without" in kwargs else []
-                os.environ['DEEP_LEARNING_API_CONFIG_PATH'] = key_tmp
+                os.environ['KONFAI_CONFIG_PATH'] = key_tmp
                 with Config(filename, key_tmp) as config:
-                    os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "False"
+                    os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
                     kwargs = {} 
                     for param in list(inspect.signature(function).parameters.values())[len(args):]:
+
                         annotation = param.annotation
+                        # --- support Literal ---
+                        if get_origin(annotation) is Literal:
+                            allowed_values = get_args(annotation)
+                            default_value = param.default if param.default != inspect._empty else allowed_values[0]
+                            value = config.getValue(param.name, f"default:{default_value}")
+                            if value not in allowed_values:
+                                raise ValueError(
+                                    f"[Config] Invalid value '{value}' for parameter '{param.name}'. "
+                                    f"Expected one of: {allowed_values}."
+                                )
+                            kwargs[param.name] = value
+                            continue
                         if str(annotation).startswith("typing.Union") or str(annotation).startswith("typing.Optional"):
                             for i in annotation.__args__:
                                 annotation = i
@@ -188,8 +201,24 @@ def config(key : Union[str, None] = None):
                             continue
                         if not annotation == inspect._empty:
                             if annotation not in [int, str, bool, float, torch.Tensor]:
-                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple"):
-                                    if annotation.__args__[0] in [int, str, bool, float]:
+                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple") or str(annotation).startswith("typing.List"):
+                                    elem_type = annotation.__args__[0]
+                                    values = config.getValue(param.name, param.default)
+                                    if getattr(elem_type, '__origin__', None) is Union:
+                                        valid_types = elem_type.__args__
+                                        result = []
+                                        for v in values:
+                                            for t in valid_types:
+                                                try:
+                                                    if t == torch.Tensor and not isinstance(v, torch.Tensor):
+                                                        v = torch.tensor(v)
+                                                    result.append(t(v) if t != torch.Tensor else v)
+                                                    break
+                                                except Exception:
+                                                    continue
+                                        kwargs[param.name] = result
+
+                                    elif annotation.__args__[0] in [int, str, bool, float]:
                                         values = config.getValue(param.name, param.default)
                                         kwargs[param.name] = values
                                     else:
@@ -204,9 +233,13 @@ def config(key : Union[str, None] = None):
                                     else: 
                                         raise ConfigError()
                                 else:
-                                    kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
-                                    if os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] == "True":
-                                        os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "False"
+                                    try:
+                                        kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
+                                    except Exception as e:
+                                        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"
                                         kwargs[param.name] = None
                             else:
                                 kwargs[param.name] = config.getValue(param.name, param.default)

konfai/utils/dataset.py

@@ -750,7 +750,7 @@ class Dataset():
         else:
             with Dataset.File(self.filename, True, self.format) as file:
                 names = file.getNames(groups)
-        return [name for i, name in enumerate(names) if index is None or i in index]
+        return [name for i, name in enumerate(sorted(names)) if index is None or i in index]
     
     def getInfos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
         if self.is_directory: