konfai 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

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, MeasureError
+from konfai.utils.utils import State, _getModule, getDevice, getGPUMemory, MeasureError, TrainerError
 from konfai.data.patching import Accumulator, ModelPatch
 
 class NetState(Enum):
@@ -41,50 +41,41 @@ class OptimizerLoader():
     def getOptimizer(self, key: str, parameter: Iterator[torch.nn.parameter.Parameter]) -> torch.optim.Optimizer:
         return config("{}.Model.{}.Optimizer".format(KONFAI_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
         
-class SchedulerStep():
-    
-    @config(None)
-    def __init__(self, nb_step : int = 0) -> None:
-        self.nb_step = nb_step
 
 class LRSchedulersLoader():
         
-    @config("Schedulers")
-    def __init__(self, params: dict[str, SchedulerStep] = {"default:ReduceLROnPlateau" : SchedulerStep(0)}) -> None:
-        self.params = params
-
-    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:
-                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():
+    @config()
+    def __init__(self, nb_step : int = 0) -> None:
+        self.nb_step = nb_step
+
+    def getschedulers(self, key: str, scheduler_classname: str, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
+        for m in ["torch.optim.lr_scheduler", "konfai.metric.schedulers"]:
+            module, name = _getModule(scheduler_classname, m)
+            if hasattr(importlib.import_module(module), name):
+                return config("{}.Model.{}.schedulers.{}".format(KONFAI_ROOT(), key, scheduler_classname))(getattr(importlib.import_module(module), name))(optimizer, config = None)
+        raise TrainerError("Unknown scheduler {}, tried importing from: 'torch.optim.lr_scheduler' and 'konfai.metric.schedulers', but no valid match was found. Check your YAML config or scheduler name spelling.".format(scheduler_classname))
+
+class LossSchedulersLoader():
         
-    @config("Schedulers")
-    def __init__(self, params: dict[str, SchedulerStep] = {"default:Constant" : SchedulerStep(0)}) -> None:
-        self.params = params
-
-    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:    
-                schedulers[getattr(importlib.import_module("konfai.metric.schedulers"), name)(config = None, DL_args = key)] = step.nb_step
-        return schedulers
+    @config()
+    def __init__(self, nb_step : int = 0) -> None:
+        self.nb_step = nb_step
+
+    def getschedulers(self, key: str, scheduler_classname: str) -> torch.optim.lr_scheduler._LRScheduler:
+        return config("{}.{}".format(key, scheduler_classname))(getattr(importlib.import_module('konfai.metric.schedulers'), scheduler_classname))(config = None)
     
 class CriterionsAttr():
     
     @config()
-    def __init__(self, l: SchedulersLoader = SchedulersLoader(), isLoss: bool = True, group: int = 0, stepStart: int = 0, stepStop: Union[int, None] = None, accumulation: bool = False) -> None:
-        self.l = l
+    def __init__(self, schedulers: dict[str, LossSchedulersLoader] = {"default:Constant": LossSchedulersLoader(0)}, isLoss: bool = True, group: int = 0, stepStart: int = 0, stepStop: Union[int, None] = None, accumulation: bool = False) -> None:
+        self.schedulersLoader = schedulers
         self.isTorchCriterion = True
         self.isLoss = isLoss
         self.stepStart = stepStart
         self.stepStop = stepStop
         self.group = group
         self.accumulation = accumulation
-        self.sheduler = None
+        self.schedulers: dict[Scheduler, int] = {}
         
 class CriterionsLoader():
 
@@ -95,9 +86,10 @@ class CriterionsLoader():
     def getCriterions(self, model_classname : str, output_group : str, target_group : str) -> dict[torch.nn.Module, CriterionsAttr]:
         criterions = {}
         for module_classpath, criterionsAttr in self.criterionsLoader.items():
-            module, name = _getModule(module_classpath, "metric.measure")
+            module, name = _getModule(module_classpath, "konfai.metric.measure")
             criterionsAttr.isTorchCriterion = module.startswith("torch")
-            criterionsAttr.sheduler = criterionsAttr.l.getschedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))
+            for scheduler_classname, schedulers in criterionsAttr.schedulersLoader.items():
+                criterionsAttr.schedulers[schedulers.getschedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}.schedulers".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath), scheduler_classname)] = schedulers.nb_step
             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
 
@@ -195,7 +187,7 @@ class Measure():
 
             for criterion, criterionsAttr in self.outputsCriterions[output_group][target_group].items():
                 if it >= criterionsAttr.stepStart and (criterionsAttr.stepStop is None or it <= criterionsAttr.stepStop):
-                    scheduler = self.update_scheduler(criterionsAttr.sheduler, it)
+                    scheduler = self.update_scheduler(criterionsAttr.schedulers, it)
                     self._loss[criterionsAttr.group]["{}:{}:{}".format(output_group, target_group, criterion.__class__.__name__)].add(scheduler.get_value(), criterion(output, *target))
                     if training and len(np.unique([len(l) for l in self._loss[criterionsAttr.group].values() if l.accumulation and l.isLoss])) == 1:
                         if criterionsAttr.isLoss:
@@ -535,7 +527,7 @@ class Network(ModuleArgsDict, ABC):
     def __init__(   self,
                     in_channels : int = 1,
                     optimizer: Union[OptimizerLoader, None] = None, 
-                    schedulers: Union[LRSchedulersLoader, None] = None, 
+                    schedulers: Union[dict[str, LRSchedulersLoader], None] = None, 
                     outputsCriterions: Union[dict[str, TargetCriterionsLoader], None] = None,
                     patch : Union[ModelPatch, None] = None,
                     nb_batch_per_step : int = 1,
@@ -549,7 +541,7 @@ class Network(ModuleArgsDict, ABC):
         self.optimizer : Union[torch.optim.Optimizer, None] = None
 
         self.LRSchedulersLoader  = schedulers
-        self.schedulers : Union[dict[torch.optim.lr_scheduler._LRScheduler, int], None] = None
+        self.schedulers : Union[dict[torch.optim.lr_scheduler._LRScheduler, int], None] = {}
 
         self.outputsCriterionsLoader = outputsCriterions
         self.measure : Union[Measure, None] = None
@@ -561,6 +553,7 @@ class Network(ModuleArgsDict, ABC):
         self.init_gain  = init_gain
         self.dim = dim
         self._it = 0
+        self._nb_lr_update = 0
         self.outputsGroup : list[OutputsGroup]= []
 
     @_function_network()
@@ -655,9 +648,15 @@ class Network(ModuleArgsDict, ABC):
                 else:
                     model_state_dict[alias] = model_state_dict_tmp[alias]
             self.load_state_dict(model_state_dict)
-            
-        if "{}_optimizer_state_dict".format(name) in state_dict and self.optimizer:
-            self.optimizer.load_state_dict(state_dict['{}_optimizer_state_dict'.format(name)])
+        if "{}_optimizer_state_dict".format(self.getName()) in state_dict and self.optimizer:
+            self.optimizer.load_state_dict(state_dict['{}_optimizer_state_dict'.format(self.getName())])
+        if "{}_it".format(self.getName()) in state_dict:
+            self._it = int(state_dict["{}_it".format(self.getName())])
+        if "{}_nb_lr_update".format(self.getName()) in state_dict:
+            self._nb_lr_update = int(state_dict["{}_nb_lr_update".format(self.getName())])
+        for scheduler in self.schedulers:
+            if scheduler.last_epoch == -1:
+                scheduler.last_epoch = self._nb_lr_update
         self.initialized()
 
     def _compute_channels_trace(self, module : ModuleArgsDict, in_channels : int, gradient_checkpoints: Union[list[str], None], gpu_checkpoints: Union[list[str], None], name: Union[str, None] = None, in_is_channel = True, out_channels : Union[int, None] = None, out_is_channel = True) -> tuple[int, bool, int, bool]:
@@ -730,7 +729,8 @@ class Network(ModuleArgsDict, ABC):
                 self.optimizer.zero_grad()
 
             if self.LRSchedulersLoader and self.optimizer:
-                self.schedulers = self.LRSchedulersLoader.getschedulers(key, self.optimizer)
+                for schedulers_classname, schedulers in self.LRSchedulersLoader.items():
+                    self.schedulers[schedulers.getschedulers(key, schedulers_classname, self.optimizer)] = schedulers.nb_step
     
     def initialized(self):
         pass
@@ -878,6 +878,7 @@ class Network(ModuleArgsDict, ABC):
                 model._requires_grad(list(self.measure.outputsCriterions.keys()))
                 for loss in self.measure.getLoss():
                     self.scaler.scale(loss / self.nb_batch_per_step).backward()
+                    
                     if self._it % self.nb_batch_per_step == 0:
                         self.scaler.step(self.optimizer)
                         self.scaler.update()
@@ -886,19 +887,19 @@ class Network(ModuleArgsDict, ABC):
 
     @_function_network()
     def update_lr(self):
+        self._nb_lr_update+=1 
         step = 0
         scheduler = None
-        if self.schedulers:
-            for scheduler, value in self.schedulers.items():
-                if value is None or (self._it >= step  and self._it < step+value):
-                    break
-                step += value
+        for scheduler, value in self.schedulers.items():
+            if value is None or (self._nb_lr_update >= step and self._nb_lr_update < step+value):
+                break
+            step += value
         if scheduler:
             if scheduler.__class__.__name__ == 'ReduceLROnPlateau':
                 if self.measure:
                     scheduler.step(sum(self.measure.getLastValues(0).values()))
             else:
-                scheduler.step()     
+                scheduler.step()
 
     @_function_network()
     def getNetworks(self) -> Self:
@@ -916,6 +917,12 @@ class Network(ModuleArgsDict, ABC):
                     v = Network.to(v, int(os.environ["device"]))
                 else:
                     v = v.to(getDevice(int(os.environ["device"])))
+        if isinstance(module, Network):
+            if module.optimizer is not None:
+                for state in module.optimizer.state.values():
+                    for k, v in state.items():
+                        if isinstance(v, torch.Tensor):
+                            state[k] = v.to(getDevice(int(os.environ["device"])))
         return module
                 
     def getName(self) -> str:
@@ -944,7 +951,7 @@ class ModelLoader():
 
     @config("Model")
     def __init__(self, classpath : str = "default:segmentation.UNet.UNet") -> None:
-        self.module, self.name = _getModule(classpath, "models")
+        self.module, self.name = _getModule(classpath, "konfai.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:

konfai/predictor.py

@@ -27,17 +27,19 @@ from collections import defaultdict
 
 class OutDataset(Dataset, NeedDevice, ABC):
 
-    def __init__(self, filename: str, group: str, pre_transforms : dict[str, TransformLoader], post_transforms : dict[str, TransformLoader], final_transforms : dict[str, TransformLoader], patchCombine: Union[str, None]) -> None: 
+    def __init__(self, filename: str, group: str, before_reduction_transforms : dict[str, TransformLoader], after_reduction_transforms : dict[str, TransformLoader], final_transforms : dict[str, TransformLoader], patchCombine: Union[str, None], reduction: str) -> None: 
         filename, format = filename.split(":")
         super().__init__(filename, format)
         self.group = group
-        self._pre_transforms = pre_transforms
-        self._post_transforms = post_transforms
+        self._before_reduction_transforms = before_reduction_transforms
+        self._after_reduction_transforms = after_reduction_transforms
         self._final_transforms = final_transforms
         self._patchCombine = patchCombine
+        self.reduction_classpath = reduction
+        self.reduction = None
        
-        self.pre_transforms : list[Transform] = []
-        self.post_transforms : list[Transform] = []
+        self.before_reduction_transforms : list[Transform] = []
+        self.after_reduction_transforms : list[Transform] = []
         self.final_transforms : list[Transform] = []
         self.patchCombine: PathCombine = None
 
@@ -47,7 +49,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
         self.nb_data_augmentation = 0
 
     def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
-        transforms_type = ["pre_transforms", "post_transforms", "final_transforms"]
+        transforms_type = ["before_reduction_transforms", "after_reduction_transforms", "final_transforms"]
         for name, _transform_type, transform_type in [(k, getattr(self, "_{}".format(k)), getattr(self, k)) for k in transforms_type]:
             
             if _transform_type is not None:
@@ -57,8 +59,15 @@ class OutDataset(Dataset, NeedDevice, ABC):
                     transform_type.append(transform)
 
         if self._patchCombine is not None:
-            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))
+            module, name = _getModule(self._patchCombine, "konfai.data.patching")                                                      
+            self.patchCombine = config("{}.outsDataset.{}.OutDataset".format(KONFAI_ROOT(), name_layer))(getattr(importlib.import_module(module), name))(config = None)
+
+        module, name = _getModule(self.reduction_classpath, "konfai.predictor")
+        if module == "konfai.predictor":
+            self.reduction = getattr(importlib.import_module(module), name)
+        else:
+            self.reduction = config("{}.outsDataset.{}.OutDataset.{}".format(KONFAI_ROOT(), name_layer, self.reduction_classpath))(getattr(importlib.import_module(module), name))(config = None)
+       
     
     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:
@@ -70,14 +79,14 @@ class OutDataset(Dataset, NeedDevice, ABC):
     
     def setDevice(self, device: torch.device):
         super().setDevice(device)
-        transforms_type = ["pre_transforms", "post_transforms", "final_transforms"]
+        transforms_type = ["before_reduction_transforms", "after_reduction_transforms", "final_transforms"]
         for transform_type in [(getattr(self, k)) for k in transforms_type]:
             if transform_type is not None:
                 for transform in transform_type:
                     transform.setDevice(device)       
 
     @abstractmethod
-    def addLayer(self, index: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
+    def addLayer(self, index_dataset: int, index_augmentation: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
         pass
 
     def isDone(self, index: int) -> bool:
@@ -96,25 +105,28 @@ class Reduction():
     def __init__(self):
         pass
 
-class ReductionMean():
+class Mean(Reduction):
 
     def __init__(self):
         pass
-    
-class ReductionMedian():
+
+    def __call__(self, result: torch.Tensor) -> torch.Tensor:
+        return torch.mean(result.float(), dim=0)
+        
+class Median(Reduction):
 
     def __init__(self):
         pass
     
-    
+    def __call__(self, result: torch.Tensor) -> torch.Tensor:
+        return torch.median(result.float(), dim=0).values
 
 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, reduction: str = "mean", inverse_transform: bool = True) -> None:
-        super().__init__(dataset_filename, group, pre_transforms, post_transforms, final_transforms, patchCombine)
+    def __init__(self, sameAsGroup: str = "default", dataset_filename: str = "default:./Dataset:mha", group: str = "default", before_reduction_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, after_reduction_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, before_reduction_transforms, after_reduction_transforms, final_transforms, patchCombine, reduction)
         self.group_src, self.group_dest = sameAsGroup.split(":")
-        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):
@@ -131,9 +143,6 @@ class OutSameAsGroupDataset(OutDataset):
             for i in range(len(input_dataset.patch.getPatch_slices(index_augmentation))):
                 self.attributes[index_dataset][index_augmentation][i] = Attribute(input_dataset.cache_attributes[0])
 
-        for transform in self.pre_transforms:
-            layer = transform(self.names[index_dataset], layer, self.attributes[index_dataset][index_augmentation][index_patch])
-
         if self.inverse_transform:
             for transform in reversed(dataset.groups_src[self.group_src][self.group_dest].post_transforms):
                 layer = transform.inverse(self.names[index_dataset], layer, self.attributes[index_dataset][index_augmentation][index_patch])
@@ -142,7 +151,7 @@ class OutSameAsGroupDataset(OutDataset):
 
     def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
         super().load(name_layer, datasets, groups)
-
+        
         if self.group_src not in groups.keys():
             raise PredictorError(
                 f"Source group '{self.group_src}' not found. Available groups: {list(groups.keys())}."
@@ -155,7 +164,6 @@ class OutSameAsGroupDataset(OutDataset):
     
     def _getOutput(self, index: int, index_augmentation: int, dataset: DatasetIter) -> torch.Tensor:
         layer = self.output_layer_accumulator[index][index_augmentation].assemble()
-        name = self.names[index]
         if index_augmentation > 0:
             
             i = 0
@@ -167,59 +175,27 @@ class OutSameAsGroupDataset(OutDataset):
                     break
                 i += dataAugmentations.nb
 
-        for transform in self.post_transforms:
-            layer = transform(name, layer, self.attributes[index][index_augmentation][0])
-            
-        if self.inverse_transform:
-            for transform in reversed(dataset.groups_src[self.group_src][self.group_dest].pre_transforms):
-                layer = transform.inverse(name, layer, self.attributes[index][index_augmentation][0])
+        for transform in self.before_reduction_transforms:
+            layer = transform(self.names[index], layer, self.attributes[index][index_augmentation][0])
+        
         return layer
 
     def getOutput(self, index: int, dataset: DatasetIter) -> torch.Tensor:
         result = torch.cat([self._getOutput(index, index_augmentation, dataset).unsqueeze(0) for index_augmentation in self.output_layer_accumulator[index].keys()], dim=0)
-        name = self.names[index]
         self.output_layer_accumulator.pop(index)
-        dtype = result.dtype
+        result = self.reduction(result.float()).to(result.dtype)
+            
+        for transform in self.after_reduction_transforms:
+            result = transform(self.names[index], result, self.attributes[index][0][0])
 
-        if self.reduction == "mean":
-            result = torch.mean(result.float(), dim=0).to(dtype)
-        elif self.reduction == "median":
-            result, _ = torch.median(result.float(), dim=0)
-        else:
-            raise NameError("Reduction method does not exist (mean, median)")
+        if self.inverse_transform:
+            for transform in reversed(dataset.groups_src[self.group_src][self.group_dest].pre_transforms):
+                result = transform.inverse(self.names[index], result, self.attributes[index][0][0])
+        
         for transform in self.final_transforms:
-            result = transform(name, result, self.attributes[index][0][0])
+            result = transform(self.names[index], result, self.attributes[index][0][0]) 
         return result
 
-class OutLayerDataset(OutDataset):
-
-    @config("OutDataset")
-    def __init__(self, dataset_filename: str = "Dataset.h5", group: str = "default", overlap : Union[list[int], None] = None, 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) -> None:
-        super().__init__(dataset_filename, group, pre_transforms, post_transforms, final_transforms, patchCombine)
-        self.overlap = overlap
-        
-    def addLayer(self, index: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
-        if index not in self.output_layer_accumulator:
-            group = list(dataset.groups.keys())[0]
-            patch_slices = get_patch_slices_from_nb_patch_per_dim(list(layer.shape[2:]), dataset.getDatasetFromIndex(group, index).patch.nb_patch_per_dim, self.overlap)
-            self.output_layer_accumulator[index] = Accumulator(patch_slices, self.patchCombine, batch=False)
-            self.attributes[index] = Attribute()
-            self.names[index] = dataset.getDatasetFromIndex(group, index).name
-            
-
-        for transform in self.pre_transforms:
-            layer = transform(layer, self.attributes[index])
-        self.output_layer_accumulator[index].addLayer(index_patch, layer)
-
-    def getOutput(self, index: int, dataset: DatasetIter) -> torch.Tensor:
-        layer = self.output_layer_accumulator[index].assemble()
-        name = self.names[index]
-        for transform in self.post_transforms:
-            layer = transform(name, layer, self.attributes[index])
-
-        self.output_layer_accumulator.pop(index)
-        return layer
-
 class OutDatasetLoader():
 
     @config("OutDataset")
@@ -307,9 +283,9 @@ class _Predictor():
 
 class ModelComposite(Network):
 
-    def __init__(self, model: Network, nb_models: int, method: str):
+    def __init__(self, model: Network, nb_models: int, combine: Reduction):
         super().__init__(model.in_channels, model.optimizer, model.schedulers, model.outputsCriterionsLoader, model.patch, model.nb_batch_per_step, model.init_type, model.init_gain, model.dim)
-        self.method = method
+        self.combine = combine
         for i in range(nb_models):
             self.add_module("Model_{}".format(i), copy.deepcopy(model), in_branch=[0], out_branch=["output_{}".format(i)])
 
@@ -330,12 +306,7 @@ class ModelComposite(Network):
 
         final_outputs = []
         for key, tensors in aggregated.items():
-            stacked = torch.stack(tensors, dim=0)
-            if self.method == 'mean':
-                agg = torch.mean(stacked, dim=0)
-            elif self.method == 'median':
-                agg = torch.median(stacked, dim=0).values
-            final_outputs.append((key, agg))
+            final_outputs.append((key, self.combine(torch.stack(tensors, dim=0))))
 
         return final_outputs
     
@@ -358,7 +329,7 @@ class Predictor(DistributedObject):
         super().__init__(train_name)
         self.manual_seed = manual_seed
         self.dataset = dataset
-        self.combine = combine
+        self.combine_classpath = combine
         self.autocast = autocast
 
         self.model = model.getModel(train=False)
@@ -430,7 +401,15 @@ class Predictor(DistributedObject):
                     "Available modules: {}".format(modules),
                     "Please check that the name matches exactly a submodule or output of your model architecture."
                 )
-        self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
+        
+        module, name = _getModule(self.combine_classpath, "konfai.predictor")
+        if module == "konfai.predictor":
+            combine = getattr(importlib.import_module(module), name)
+        else:
+            combine = config("{}.{}".format(KONFAI_ROOT(), self.combine_classpath))(getattr(importlib.import_module(module), name))(config = None)
+       
+        
+        self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), combine)
         self.modelComposite.load(self._load())
 
         if len(list(self.outsDataset.keys())) == 0 and len([network for network in self.modelComposite.getNetworks().values() if network.measure is not None]) == 0:

konfai/trainer.py

@@ -108,9 +108,7 @@ class _Trainer():
         if data_log is not None:
             for data in data_log:
                 self.data_log[data.split("/")[0].replace(":", ".")] = (DataLog.__getitem__(data.split("/")[1]).value[0], int(data.split("/")[2]))
-        
-       
-        
+ 
     def __enter__(self):
         return self
     
@@ -212,6 +210,9 @@ class _Trainer():
             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})
+        save_dict.update({'{}_it'.format(name): network._it for name, network in self.model.module.getNetworks().items() if network.optimizer is not None})
+        save_dict.update({'{}_nb_lr_update'.format(name): network._nb_lr_update 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":
@@ -356,7 +357,7 @@ class Trainer(DistributedObject):
                     name = sorted(os.listdir(path))[-1]
             
             if os.path.exists(path+name):
-                state_dict = torch.load(path+name, weights_only=False)
+                state_dict = torch.load(path+name, weights_only=False, map_location="cpu")
             else:
                 raise Exception("Model : {} does not exist !".format(self.name))
             

konfai/utils/utils.py

@@ -20,6 +20,11 @@ import torch.nn.functional as F
 import sys
 import re
 
+import requests
+from tqdm import tqdm
+import importlib
+from pathlib import Path
+import shutil
 
 def description(model, modelEMA = None, showMemory: bool = True, train: bool = True) -> str:
     values_desc = lambda weights, values: " ".join(["{}({:.2f}) : {:.6f}".format(name.split(":")[-1], weight, value) for (name, value), weight in zip(values.items(), weights.values())])
@@ -35,11 +40,11 @@ def description(model, modelEMA = None, showMemory: bool = True, train: bool = T
 def _getModule(classpath : str, type : str) -> tuple[str, str]:
     if len(classpath.split(":")) > 1:
         module = ".".join(classpath.split(":")[:-1])
-        name = classpath.split(":")[-1] 
+        name = classpath.split(":")[-1]
     else:
-        module = "konfai."+type+("." if len(classpath.split(".")) > 2 else "")+".".join(classpath.split(".")[:-1])
+        module = type+("." if len(classpath.split(".")) > 2 else "")+".".join(classpath.split(".")[:-1])
         name = classpath.split(".")[-1]
-    return module, name
+    return module, name.split("/")[0]
 
 def cpuInfo() -> str:
     return "CPU ({:.2f} %)".format(psutil.cpu_percent(interval=0.5))
@@ -420,6 +425,8 @@ def setup(parser: argparse.ArgumentParser) -> DistributedObject:
 
     os.environ["CUDA_VISIBLE_DEVICES"] = config["gpu"]
     os.environ["KONFAI_NB_CORES"] = config["cpu"]
+
+    os.environ["KONFAI_WORKERS"] = str(config["num_workers"])
     os.environ["KONFAI_MODELS_DIRECTORY"] = config["MODELS_DIRECTORY"]
     os.environ["KONFAI_CHECKPOINTS_DIRECTORY"] = config["CHECKPOINTS_DIRECTORY"]
     os.environ["KONFAI_PREDICTIONS_DIRECTORY"] = config["PREDICTIONS_DIRECTORY"]
@@ -536,6 +543,37 @@ def _resample_affine(data: torch.Tensor, matrix: torch.Tensor):
     return F.grid_sample(data.unsqueeze(0).type(torch.float32), F.affine_grid(matrix[:, :-1,...].type(torch.float32), [1]+list(data.shape), align_corners=True), align_corners=True, mode=mode, padding_mode="reflection").squeeze(0).type(data.dtype)
 
 
+
+def download_url(model_name: str, url: str) -> str:
+    spec = importlib.util.find_spec("konfai")
+    base_path = Path(spec.submodule_search_locations[0]) / "metric" / "models"
+    subdirs = Path(model_name).parent
+    model_dir = base_path / subdirs
+    model_dir.mkdir(exist_ok=True)
+    filetmp = model_dir / ("tmp_"+str(Path(model_name).name))
+    file = model_dir / Path(model_name).name
+    if file.exists():
+        return str(file)
+    
+    try:
+        print(f"[FOCUS] Downloading {model_name} to {file}")
+        with requests.get(url+model_name, stream=True) as r:
+            r.raise_for_status()
+            total = int(r.headers.get('content-length', 0))
+            with open(filetmp, 'wb') as f:
+                with tqdm(total=total, unit='B', unit_scale=True, desc=f"Downloading {model_name}") as pbar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+        shutil.copy2(filetmp, file)
+        print("Download finished.")
+    except Exception as e:
+        raise e
+    finally:
+        if filetmp.exists():
+            os.remove(filetmp)
+    return str(file)
+    
 SUPPORTED_EXTENSIONS = [
     "mha", "mhd",         # MetaImage
     "nii", "nii.gz",      # NIfTI

{konfai-1.1.5.dist-info → konfai-1.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.5
+Version: 1.1.7
 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.1.7.dist-info/RECORD

@@ -0,0 +1,39 @@
+konfai/__init__.py,sha256=YXG-wpSEXWs6Jt3BDI77V4r89gEUNX-6lxW9btj5VYI,851
+konfai/evaluator.py,sha256=WM78NGydV0iqKElahr-WpOCZLEJXHjjPq-LS-gi23rk,8401
+konfai/main.py,sha256=kr7Iie_f67NF6G3dAAj9G6Z9dhn9RzbdLYpzy2WvIh8,2573
+konfai/predictor.py,sha256=6zFGjQhG5-3hF10pvCn58UDLB6ZQ_UhMlxkOCzpM4tM,23053
+konfai/trainer.py,sha256=Edi8l8OOfCewYM-Cd5C5rCqCaprvlfxzohd4iLkK5u0,20632
+konfai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+konfai/data/augmentation.py,sha256=mFVMpbJ8WBKGbMILdmTZYBA8k7kRDQVPOEy1A9t5QP4,32281
+konfai/data/data_manager.py,sha256=yphkTjk4_gyr_vOGodfhu9ImDHplRe2KR6dL2ORXzCw,29000
+konfai/data/patching.py,sha256=zAm6jjUW--lsqTBlDFICVlj7O_QOlSxStFAHS9S9H8I,14753
+konfai/data/transform.py,sha256=yZ6aALtVEqTYghREPH8Z1deglU7M4t4OiQqMY6gpjjA,26559
+konfai/metric/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+konfai/metric/measure.py,sha256=S65WGBRYMoalkwI_Le3C6K0tQmT4Qft82nUTBwhpvFU,26907
+konfai/metric/schedulers.py,sha256=VPp7zEwD0AtQz51XG0TlutD_NrsTZs4fstC5h8A8f8U,1309
+konfai/models/classification/convNeXt.py,sha256=hZQ_9I1lJW7DX8QnoeZ9L1br-98apV9VXGrsTAN25ds,9261
+konfai/models/classification/resnet.py,sha256=eBoS_zNczfa1EnECKKLr0Ow26ekR_4W5eji6p9cPfHk,7978
+konfai/models/generation/cStyleGan.py,sha256=7D8zZveDEZapFeaaDTe3wVhuDCCHLqq6cpXtl2-QaJA,8059
+konfai/models/generation/ddpm.py,sha256=jXb0eOU3i_gMHvj9pawVAWQDjMGl6SPfCc5m2Jzdrik,13175
+konfai/models/generation/diffusionGan.py,sha256=ZFKdRHvlajEPyw9_HCLo0Q92iChCtZd2dvi2nEF_tBI,33218
+konfai/models/generation/gan.py,sha256=RzpGNu8BlBVDCxFR5CCmEFYUzVBs2YrDl7trIos3ssw,7861
+konfai/models/generation/vae.py,sha256=zH8qk04z2lXDhXAVxTHiRIVMfRi3brB61lHGmEsE3kM,4675
+konfai/models/registration/registration.py,sha256=EAE3w8aic2fPWiJz0ilqrs2kCGUQD6NWvysVfHXxA_g,6329
+konfai/models/representation/representation.py,sha256=TiYcBBqZYySpwsRlnnBQh0QVW29Rcvb9GUjsqnCKKLM,2689
+konfai/models/segmentation/NestedUNet.py,sha256=hDSE7BJ17IqaiHWAD5zlVG7G_KvQzuQmVnAVsiYVE6E,10248
+konfai/models/segmentation/UNet.py,sha256=BktCRfAcCDtvGCw8wGfyZvBtT4G0Oy8teIcVgDFOurk,4078
+konfai/network/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+konfai/network/blocks.py,sha256=C_dOLmXovxANWCx7P439FzR_95Zisy2I1F7FEwwz7AE,14412
+konfai/network/network.py,sha256=LHeA7HtsVYO7BJu2_kqh23q2GIANn5ZSa4LhKMt7dJg,48642
+konfai/utils/ITK.py,sha256=OxTieDNNYHGkn7zxJsAG-6ecRG1VYMvn1dlBbBe1DOs,13955
+konfai/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+konfai/utils/config.py,sha256=f7o83ix5_oNbr2pki-Czqr-yHi-8n92ZL64nlo0XGwA,12514
+konfai/utils/dataset.py,sha256=6ZzevdhJ7e5zlXATAVwSh9O6acKXM7gYNxkMAa5DrmM,36351
+konfai/utils/registration.py,sha256=v1srEBOcgDnHrx0YtsK6bcj0yCMH7wNeaQ3wC7gEvOw,8898
+konfai/utils/utils.py,sha256=mQQ6FB0Jw7Odg5GxYeTApR6lvFhXn_iz-GVGZRngGQA,24934
+konfai-1.1.7.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+konfai-1.1.7.dist-info/METADATA,sha256=DJ2n6Rc9NF018KsJPQxFHpQmlDazUnrrP4zHtgE4_Ew,2515
+konfai-1.1.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+konfai-1.1.7.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
+konfai-1.1.7.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
+konfai-1.1.7.dist-info/RECORD,,