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

konfai/evaluator.py

@@ -9,7 +9,7 @@ import builtins
 import importlib
 from konfai import EVALUATIONS_DIRECTORY, PREDICTIONS_DIRECTORY, KONFAI_ROOT, CONFIG_FILE
 from konfai.utils.config import config
-from konfai.utils.utils import _getModule, DistributedObject, synchronize_data
+from konfai.utils.utils import _getModule, DistributedObject, synchronize_data, EvaluatorError
 from konfai.data.data_manager import DataMetric
 
 class CriterionsAttr():
@@ -54,7 +54,8 @@ class Statistics():
             if name_dataset not in self.measures:
                 self.measures[name_dataset] = {}
             self.measures[name_dataset][name] = value
-            
+    
+    @staticmethod   
     def getStatistic(values: list[float]) -> dict[str, float]:
         return {"max": np.max(values), "min": np.min(values), "std": np.std(values), "25pc": np.percentile(values, 25), "50pc": np.percentile(values, 50), "75pc": np.percentile(values, 75), "mean": np.mean(values), "count": len(values)}
     
@@ -91,7 +92,6 @@ class Evaluator(DistributedObject):
             exit(0)
         super().__init__(train_name)
         self.metric_path = EVALUATIONS_DIRECTORY()+self.name+"/"
-        self.predict_path = PREDICTIONS_DIRECTORY()+self.name+"/"
         self.metricsLoader = metrics
         self.dataset = dataset
         self.metrics = {k: v.getTargetsCriterions(k) for k, v in self.metricsLoader.items()}
@@ -102,10 +102,10 @@ class Evaluator(DistributedObject):
         result = {}
         for output_group in self.metrics:
             for target_group in self.metrics[output_group]:
-                targets = [data_dict[group][0] for group in target_group.split("/") if group in data_dict]
+                targets = [data_dict[group][0].to(0) if torch.cuda.is_available() else data_dict[group][0] for group in target_group.split("/") if group in data_dict]
                 name = data_dict[output_group][1][0]
                 for metric in self.metrics[output_group][target_group]:
-                    result["{}:{}:{}".format(output_group, target_group, metric.__class__.__name__)] = metric(data_dict[output_group][0], *targets).item()
+                    result["{}:{}:{}".format(output_group, target_group, metric.__class__.__name__)] = metric(data_dict[output_group][0].to(0) if torch.cuda.is_available() else data_dict[output_group][0], *targets).item()
         statistics.add(result, name)
         return result
     
@@ -126,9 +126,26 @@ class Evaluator(DistributedObject):
 
         self.dataloader = self.dataset.getData(world_size)
 
+        groupsDest = [group for groups in self.dataset.groups_src.values() for group in groups]
+
+        missing_outputs = set(self.metrics.keys()) - set(groupsDest)
+        if missing_outputs:
+            raise EvaluatorError(
+                f"The following metric output groups are missing from 'groupsDest': {sorted(missing_outputs)}. ",
+                f"Available groups: {sorted(groupsDest)}"
+            )
+
+        target_groups = {target for targets in self.metrics.values() for target in targets}
+        missing_targets = target_groups - set(groupsDest)
+        if missing_targets:
+            raise EvaluatorError(
+                f"The following metric target groups are missing from 'groupsDest': {sorted(missing_targets)}. ",
+                f"Available groups: {sorted(groupsDest)}"
+            )
+
     def run_process(self, world_size: int, global_rank: int, gpu: int, dataloaders: list[DataLoader]):
         description = lambda measure : "Metric TRAIN : {} ".format(" | ".join("{}: {:.2f}".format(k, v) for k, v in measure.items()) if measure is not None else "")
-        with tqdm.tqdm(iterable = enumerate(dataloaders[0]), leave=False, desc = description(None), total=len(dataloaders[0])) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(dataloaders[0]), leave=True, desc = description(None), total=len(dataloaders[0]), ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 batch_iter.set_description(description(self.update({k: (v[0], v[4]) for k,v in data_dict.items()}, self.statistics_train)))
         outputs = synchronize_data(world_size, gpu, self.statistics_train.measures)
@@ -136,7 +153,7 @@ class Evaluator(DistributedObject):
             self.statistics_train.write(outputs)
         if len(dataloaders) == 2:
             description = lambda measure : "Metric VALIDATION : {} ".format(" | ".join("{}: {:.2f}".format(k, v) for k, v in measure.items()) if measure is not None else "")
-            with tqdm.tqdm(iterable = enumerate(dataloaders[1]), leave=False, desc = description(None), total=len(dataloaders[1])) as batch_iter:
+            with tqdm.tqdm(iterable = enumerate(dataloaders[1]), leave=True, desc = description(None), total=len(dataloaders[1]), ncols=0) as batch_iter:
                 for _, data_dict in batch_iter:
                     batch_iter.set_description(description(self.update({k: (v[0], v[4]) for k,v in data_dict.items()}, self.statistics_validation)))
             outputs = synchronize_data(world_size, gpu, self.statistics_validation.measures)

konfai/main.py

@@ -3,6 +3,7 @@ import os
 from torch.cuda import device_count
 import torch.multiprocessing as mp
 from konfai.utils.utils import setup, TensorBoard, Log
+from konfai import KONFAI_NB_CORES
 
 import sys
 sys.path.insert(0, os.getcwd())
@@ -11,16 +12,15 @@ def main():
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     try:
         with setup(parser) as distributedObject:
-            with Log(distributedObject.name):
+            with Log(distributedObject.name, 0):
                 world_size = device_count()
                 if world_size == 0:
-                    world_size = 1
+                    world_size = int(KONFAI_NB_CORES())
                 distributedObject.setup(world_size)
                 with TensorBoard(distributedObject.name):
                     mp.spawn(distributedObject, nprocs=world_size)
-    except Exception as e:
-        print(e)
-        exit(1)
+    except KeyboardInterrupt:
+        print("\n[KonfAI] Manual interruption (Ctrl+C)")
 
 
 def cluster():
@@ -47,6 +47,8 @@ def cluster():
             executor.update_parameters(name=config["name"], mem_gb=config["memory"], gpus_per_node=n_gpu, tasks_per_node=n_gpu//distributedObject.size, cpus_per_task=config["num_workers"], nodes=config["num_nodes"], timeout_min=config["time_limit"])
             with TensorBoard(distributedObject.name):
                 executor.submit(distributedObject)
+    except KeyboardInterrupt:
+        print("\n[KonfAI] Manual interruption (Ctrl+C)")
     except Exception as e:
         print(e)
         exit(1)

konfai/models/registration/registration.py

@@ -94,7 +94,6 @@ class SpatialTransformer(torch.nn.Module):
             new_locs[:, 1,1] = 1
             new_locs[:, 0,2] = flow[:, 0]
             new_locs[:, 1,2] = flow[:, 1]
-            print(new_locs)
             return F.grid_sample(src, F.affine_grid(new_locs, src.size()), align_corners=True, mode="bilinear")
         else:
             new_locs = self.grid + flow

konfai/network/blocks.py

@@ -176,7 +176,6 @@ class Print(torch.nn.Module):
         super().__init__()
     
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        print(input.shape)
         return input
 
 class Write(torch.nn.Module):

konfai/network/network.py

@@ -16,7 +16,7 @@ from enum import Enum
 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.utils.utils import State, _getModule, getDevice, getGPUMemory, MeasureError
 from konfai.data.patching import Accumulator, ModelPatch
 
 class NetState(Enum):
@@ -39,7 +39,6 @@ class OptimizerLoader():
         self.name = name
     
     def getOptimizer(self, key: str, parameter: Iterator[torch.nn.parameter.Parameter]) -> torch.optim.Optimizer:
-        torch.optim.AdamW
         return config("{}.Model.{}.Optimizer".format(KONFAI_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
         
 class SchedulerStep():
@@ -54,12 +53,12 @@ class LRSchedulersLoader():
     def __init__(self, params: dict[str, SchedulerStep] = {"default:ReduceLROnPlateau" : SchedulerStep(0)}) -> None:
         self.params = params
 
-    def getShedulers(self, key: str, optimizer: torch.optim.Optimizer) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
-        shedulers : dict[torch.optim.lr_scheduler._LRScheduler, int] = {}
+    def getschedulers(self, key: str, optimizer: torch.optim.Optimizer) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
+        schedulers : dict[torch.optim.lr_scheduler._LRScheduler, int] = {}
         for name, step in self.params.items():
             if name:
-                shedulers[config("Trainer.Model.{}.Schedulers.{}".format(key, name))(getattr(importlib.import_module('torch.optim.lr_scheduler'), name))(optimizer, config = None)] = step.nb_step
-        return shedulers
+                schedulers[config("Trainer.Model.{}.Schedulers.{}".format(key, name))(getattr(importlib.import_module('torch.optim.lr_scheduler'), name))(optimizer, config = None)] = step.nb_step
+        return schedulers
 
 class SchedulersLoader():
         
@@ -67,12 +66,12 @@ class SchedulersLoader():
     def __init__(self, params: dict[str, SchedulerStep] = {"default:Constant" : SchedulerStep(0)}) -> None:
         self.params = params
 
-    def getShedulers(self, key: str) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
-        shedulers : dict[Scheduler, int] = {}
+    def getschedulers(self, key: str) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
+        schedulers : dict[Scheduler, int] = {}
         for name, step in self.params.items():
             if name:    
-                shedulers[getattr(importlib.import_module("konfai.metric.schedulers"), name)(config = None, DL_args = key)] = step.nb_step
-        return shedulers
+                schedulers[getattr(importlib.import_module("konfai.metric.schedulers"), name)(config = None, DL_args = key)] = step.nb_step
+        return schedulers
     
 class CriterionsAttr():
     
@@ -98,7 +97,7 @@ 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(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))
+            criterionsAttr.sheduler = criterionsAttr.l.getschedulers("{}.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
 
@@ -154,10 +153,25 @@ class Measure():
             self.outputsCriterions[output_group.replace(":", ".")] = targetCriterionsLoader.getTargetsCriterions(output_group, model_classname)
         self._loss : dict[int, dict[str, Measure.Loss]] = {}
 
-    def init(self, model : torch.nn.Module) -> None:
+    def init(self, model : torch.nn.Module, group_dest: list[str]) -> None:
         outputs_group_rename = {}
+        
+        modules = []
+        for i,_ in model.named_modules():
+            modules.append(i)
+
         for output_group in self.outputsCriterions.keys():
+            if output_group not in modules:
+                  raise MeasureError(f"The output group '{output_group}' defined in 'outputsCriterions' does not correspond to any module in the model.",
+                    f"Available modules: {modules}",
+                    "Please check that the name matches exactly a submodule or output of your model architecture."
+                )
             for target_group in self.outputsCriterions[output_group]:
+                if target_group not in group_dest:
+                    raise MeasureError(
+                        f"The target_group '{target_group}' defined in 'outputsCriterions.{output_group}.targetsCriterions' was not found in the available destination groups.",
+                        "This target_group is expected for loss or metric computation, but was not loaded in 'group_dest'.",
+                        f"Please make sure that the group '{target_group}' is defined in 'Dataset:groups_src:...:groups_dest:'{target_group}'' and correctly loaded from the dataset.")
                 for criterion in self.outputsCriterions[output_group][target_group]:
                     if not self.outputsCriterions[output_group][target_group][criterion].isTorchCriterion:
                         outputs_group_rename[output_group] = criterion.init(model, output_group, target_group)
@@ -703,10 +717,10 @@ class Network(ModuleArgsDict, ABC):
         return in_channels, in_is_channel, out_channels, out_is_channel
 
     @_function_network()
-    def init(self, autocast : bool, state : State, key: str) -> None: 
+    def init(self, autocast : bool, state : State, group_dest: list[str], key: str) -> None: 
         if self.outputsCriterionsLoader:
             self.measure = Measure(key, self.outputsCriterionsLoader)
-            self.measure.init(self)
+            self.measure.init(self, group_dest)
         if state != State.PREDICTION:
             self.scaler = torch.amp.GradScaler("cuda", enabled=autocast)
             if self.optimizerLoader:
@@ -714,7 +728,7 @@ class Network(ModuleArgsDict, ABC):
                 self.optimizer.zero_grad()
 
             if self.LRSchedulersLoader and self.optimizer:
-                self.schedulers = self.LRSchedulersLoader.getShedulers(key, self.optimizer)
+                self.schedulers = self.LRSchedulersLoader.getschedulers(key, self.optimizer)
     
     def initialized(self):
         pass
@@ -880,7 +894,6 @@ class Network(ModuleArgsDict, ABC):
         if scheduler:
             if scheduler.__class__.__name__ == 'ReduceLROnPlateau':
                 if self.measure:
-                    print(sum(self.measure.getLastValues(0).values()))
                     scheduler.step(sum(self.measure.getLastValues(0).values()))
             else:
                 scheduler.step()     

konfai/predictor.py

@@ -94,10 +94,10 @@ class OutDataset(Dataset, NeedDevice, ABC):
 class OutSameAsGroupDataset(OutDataset):
 
     @config("OutDataset")
-    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:
+    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, reduction: 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
+        self.reduction = reduction
         self.inverse_transform = inverse_transform
 
     def addLayer(self, index_dataset: int, index_augmentation: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
@@ -151,9 +151,9 @@ class OutSameAsGroupDataset(OutDataset):
         self.output_layer_accumulator.pop(index)
         dtype = result.dtype
 
-        if self.redution == "mean":
+        if self.reduction == "mean":
             result = torch.mean(result.float(), dim=0).to(dtype)
-        elif self.redution == "median":
+        elif self.reduction == "median":
             result, _ = torch.median(result.float(), dim=0)
         else:
             raise NameError("Reduction method does not exist (mean, median)")
@@ -201,7 +201,7 @@ class OutDatasetLoader():
 
 class _Predictor():
 
-    def __init__(self, world_size: int, global_rank: int, local_rank: int, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], modelComposite: DDP, dataloader_prediction: DataLoader) -> None:
+    def __init__(self, world_size: int, global_rank: int, local_rank: int, autocast: bool, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], modelComposite: DDP, dataloader_prediction: DataLoader) -> None:
         self.world_size = world_size        
         self.global_rank = global_rank
         self.local_rank = local_rank
@@ -209,7 +209,7 @@ class _Predictor():
         self.modelComposite = modelComposite
         self.dataloader_prediction = dataloader_prediction
         self.outsDataset = outsDataset
-
+        self.autocast = autocast
         
         self.it = 0
 
@@ -239,21 +239,22 @@ class _Predictor():
         self.modelComposite.eval()  
         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 "KONFAI_CLUSTER" not in os.environ) as batch_iter:
+        self.dataloader_prediction.dataset.load("Prediction")
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=True, desc = desc(), total=len(self.dataloader_prediction), ncols=0) as batch_iter:
             dist.barrier()
             for it, data_dict in batch_iter:
-                input = self.getInput(data_dict)
-                for name, output in self.modelComposite(input, list(self.outsDataset.keys())):
-                    self._predict_log(data_dict)
-                    outDataset = self.outsDataset[name]
-                    for i, (index, patch_augmentation, patch_index) in enumerate([(int(index), int(patch_augmentation), int(patch_index)) for index, patch_augmentation, patch_index in zip(list(data_dict.values())[0][1], list(data_dict.values())[0][2], list(data_dict.values())[0][3])]):
-                        outDataset.addLayer(index, patch_augmentation, patch_index, output[i].cpu(), self.dataset)
-                        if outDataset.isDone(index):
-                            outDataset.write(index, self.dataset.getDatasetFromIndex(list(data_dict.keys())[0], index).name.split("/")[-1], outDataset.getOutput(index, self.dataset))
-
-                batch_iter.set_description(desc())
-                self.it += 1
+                with torch.amp.autocast('cuda', enabled=self.autocast):
+                    input = self.getInput(data_dict)
+                    for name, output in self.modelComposite(input, list(self.outsDataset.keys())):
+                        self._predict_log(data_dict)
+                        outDataset = self.outsDataset[name]
+                        for i, (index, patch_augmentation, patch_index) in enumerate([(int(index), int(patch_augmentation), int(patch_index)) for index, patch_augmentation, patch_index in zip(list(data_dict.values())[0][1], list(data_dict.values())[0][2], list(data_dict.values())[0][3])]):
+                            outDataset.addLayer(index, patch_augmentation, patch_index, output[i].cpu(), self.dataset)
+                            if outDataset.isDone(index):
+                                outDataset.write(index, self.dataset.getDatasetFromIndex(list(data_dict.keys())[0], index).name.split("/")[-1], outDataset.getOutput(index, self.dataset))
+
+                    batch_iter.set_description(desc())
+                    self.it += 1
 
     def _predict_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]):
         measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank, {"" : self.modelComposite.module}, 1)
@@ -320,6 +321,7 @@ class Predictor(DistributedObject):
                     train_name: str = "name",
                     manual_seed : Union[int, None] = None,
                     gpu_checkpoints: Union[list[str], None] = None,
+                    autocast : bool = False,
                     outsDataset: Union[dict[str, OutDatasetLoader], None] = {"default:Default" : OutDatasetLoader()},
                     images_log: list[str] = []) -> None:
         if os.environ["KONFAI_CONFIG_MODE"] != "Done":
@@ -328,6 +330,7 @@ class Predictor(DistributedObject):
         self.manual_seed = manual_seed
         self.dataset = dataset
         self.combine = combine
+        self.autocast = autocast
 
         self.model = model.getModel(train=False)
         self.it = 0
@@ -384,7 +387,7 @@ class Predictor(DistributedObject):
 
         shutil.copyfile(CONFIG_FILE(), self.predict_path+"Prediction.yml")
 
-        self.model.init(autocast=False, state = State.PREDICTION)
+        self.model.init(self.autocast, State.PREDICTION, self.dataset.getGroupsDest())
         self.model.init_outputsGroup()
         self.model._compute_channels_trace(self.model, self.model.in_channels, None, self.gpu_checkpoints)
         self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
@@ -402,7 +405,7 @@ class Predictor(DistributedObject):
     def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):
         modelComposite = Network.to(self.modelComposite, local_rank*self.size)
         modelComposite = DDP(modelComposite, static_graph=True) if torch.cuda.is_available() else CPU_Model(modelComposite)
-        with _Predictor(world_size, global_rank, local_rank, self.predict_path, self.images_log, self.outsDataset, modelComposite, *dataloaders) as p:
+        with _Predictor(world_size, global_rank, local_rank, self.autocast, self.predict_path, self.images_log, self.outsDataset, modelComposite, *dataloaders) as p:
             p.run()
 
         

konfai/trainer.py

@@ -14,10 +14,9 @@ import torch.distributed as dist
 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.utils.utils import State, DataLog, DistributedObject, description, TrainerError
 from konfai.network.network import Network, ModelLoader, NetState, CPU_Model
 
-
 class EarlyStoppingBase:
 
     def __init__(self):
@@ -53,8 +52,8 @@ class EarlyStopping(EarlyStoppingBase):
             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. "
+                raise TrainerError(
+                    "Metric '{}' specified in EarlyStopping.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])
         
@@ -68,7 +67,7 @@ class EarlyStopping(EarlyStoppingBase):
         elif self.mode == "max":
             improvement = current_score - self.best_score
         else:
-            raise ValueError("Mode must be 'min' or 'max'.")
+            raise TrainerError("Mode must be 'min' or 'max'.")
 
         if improvement > self.min_delta:
             self.best_score = current_score
@@ -99,7 +98,7 @@ class _Trainer():
         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)
@@ -118,13 +117,14 @@ class _Trainer():
             self.tb.close()
 
     def run(self) -> None:
-        with tqdm.tqdm(iterable = range(self.epoch, self.epochs), leave=False, total=self.epochs, initial=self.epoch, desc="Progress", disable=self.global_rank != 0) as epoch_tqdm:
+        self.dataloader_training.dataset.load("Train")
+        self.dataloader_validation.dataset.load("Validation")
+        with tqdm.tqdm(iterable = range(self.epoch, self.epochs), leave=False, total=self.epochs, initial=self.epoch, desc="Progress") as epoch_tqdm:
             for self.epoch in epoch_tqdm:
-                self.dataloader_training.dataset.load()
                 self.train()
                 if self.early_stopping.isStopped():
                     break
-                self.dataloader_training.dataset.resetAugmentation()
+                self.dataloader_training.dataset.resetAugmentation("Train")
                 
     def getInput(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int, str, bool]]) -> dict[tuple[str, bool], torch.Tensor]:
         return {(k, v[5][0].item()) : v[0] for k, v in data_dict.items()}
@@ -137,7 +137,8 @@ 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 "KONFAI_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, ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 with torch.amp.autocast('cuda', enabled=self.autocast):
                     input = self.getInput(data_dict)
@@ -171,8 +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 "KONFAI_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, ncols=0) as batch_iter:
             for _, data_dict in batch_iter:
                 input = self.getInput(data_dict)
                 self.model(input)
@@ -180,7 +180,7 @@ class _Trainer():
                     self.modelEMA.module(input)
 
                 batch_iter.set_description(desc())
-        self.dataloader_validation.dataset.resetAugmentation()
+        self.dataloader_validation.dataset.resetAugmentation("Validation")
         dist.barrier()
         self.model.train()
         self.model.module.setState(NetState.TRAIN)
@@ -360,7 +360,7 @@ class Trainer(DistributedObject):
                     os.makedirs(dir)
 
             for name in sorted(os.listdir(path_checkpoint)):
-                checkpoint = torch.load(path_checkpoint+name, weights_only=False)
+                checkpoint = torch.load(path_checkpoint+name, weights_only=False, map_location='cpu')
                 self.model.load(checkpoint, init=False, ema=False)
 
                 torch.save(self.model, "{}Serialized/{}".format(path_model, name))
@@ -391,7 +391,7 @@ class Trainer(DistributedObject):
         if state != State.TRAIN:
             state_dict = self._load()
 
-        self.model.init(self.autocast, state)
+        self.model.init(self.autocast, state, self.dataset.getGroupsDest())
         self.model.init_outputsGroup()
         self.model._compute_channels_trace(self.model, self.model.in_channels, self.gradient_checkpoints, self.gpu_checkpoints)
         self.model.load(state_dict, init=True, ema=False)

konfai/utils/config.py

@@ -6,14 +6,10 @@ from copy import deepcopy
 from typing import Union, Literal, get_origin, get_args
 import torch
 from konfai import CONFIG_FILE
+from konfai.utils.utils import ConfigError
 
 yaml = ruamel.yaml.YAML()
 
-class ConfigError(Exception):
-
-    def __init__(self, message : str = "The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]") -> None:
-        self.message = message
-        super().__init__(self.message)
 
 class Config():
 
@@ -187,9 +183,8 @@ def config(key : Union[str, None] = None):
                             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}."
+                                raise ConfigError(
+                                    f"Invalid value '{value}' for parameter '{param.name} expected one of: {allowed_values}."
                                 )
                             kwargs[param.name] = value
                             continue
@@ -222,21 +217,26 @@ def config(key : Union[str, None] = None):
                                         values = config.getValue(param.name, param.default)
                                         kwargs[param.name] = values
                                     else:
-                                        raise ConfigError()
+                                        raise ConfigError("Config: The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]")
                                 elif str(annotation).startswith("dict"):
                                     if annotation.__args__[0] == str:
                                         values = config.getValue(param.name, param.default)
                                         if values is not None and annotation.__args__[1] not in [int, str, bool, float]:
-                                            kwargs[param.name] = {value : annotation.__args__[1](config = filename, DL_args = key_tmp+"."+param.name+"."+value) for value in values}
+                                            try:
+                                                kwargs[param.name] = {value : annotation.__args__[1](config = filename, DL_args = key_tmp+"."+param.name+"."+value) for value in values}
+                                            except ValueError as e:
+                                                raise ValueError(e)
+                                            except Exception as e:
+                                                raise ConfigError("{} {}".format(values, e))
                                         else:
                                             kwargs[param.name] = values
                                     else: 
-                                        raise ConfigError()
+                                        raise ConfigError("Config: The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]")
                                 else:
                                     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))
+                                        raise ConfigError("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/utils/dataset.py

@@ -348,6 +348,10 @@ class Dataset():
         @abstractmethod
         def getNames(self, group: str) -> list[str]:
             pass
+        
+        @abstractmethod
+        def getGroup(self) -> list[str]:
+            pass
 
         @abstractmethod
         def isExist(self, group: str, name: Union[str, None] = None) -> bool:
@@ -458,6 +462,9 @@ class Dataset():
                     names.extend(self.getNames("/".join(groups.split("/")[1:]), h5_group[group]))
             return names
         
+        def getGroup(self):
+            return self.h5.keys()
+        
         def _getDataset(self, groups: str, name: str, h5_group: h5py.Group = None) -> h5py.Dataset:
             if h5_group is None:
                 h5_group = self.h5
@@ -613,7 +620,10 @@ class Dataset():
         
         def getNames(self, group: str) -> list[str]:
             raise NotImplementedError()
-
+        
+        def getGroup(self):
+            raise NotImplementedError()
+        
         def getInfos(self, group: str, name: str) -> tuple[list[int], Attribute]:
             attributes = Attribute()
             if os.path.exists("{}{}{}.{}".format(self.filename, group if group is not None else "", name, self.format)):
@@ -736,7 +746,7 @@ class Dataset():
                 subDirectories.extend(self._getSubDirectories("/".join(groups.split("/")[1:]), subDirectory))
         return subDirectories
 
-    def getNames(self, groups: str, index: Union[list[int], None] = None, subDirectory: str = "") -> list[str]:
+    def getNames(self, groups: str, index: Union[list[int], None] = None) -> list[str]:
         names = []
         if self.is_directory:
             for subDirectory in self._getSubDirectories(groups):
@@ -752,6 +762,21 @@ class Dataset():
                 names = file.getNames(groups)
         return [name for i, name in enumerate(sorted(names)) if index is None or i in index]
     
+    def getGroup(self):
+        if self.is_directory:
+            groups = set()
+            for root, _, files in os.walk(self.filename):
+                for file in files:
+                    path = os.path.relpath(os.path.join(root, file.split(".")[0]), self.filename)
+                    parts = path.split("/")
+                    if len(parts) >= 2:
+                        del parts[-2]
+                    groups.add("/".join(parts))
+        else:
+            with Dataset.File(self.filename, True, self.format) as file:
+                groups = file.getGroup()
+        return list(groups)
+    
     def getInfos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
         if self.is_directory:
             for subDirectory in self._getSubDirectories(groups):