konfai 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

konfai/data/augmentation.py

@@ -57,7 +57,7 @@ class DataAugmentationsList():
 
     def load(self, key: str, datasets: list[Dataset]):
         for augmentation, prob in self.dataAugmentationsLoader.items():
-            module, name = _getModule(augmentation, "data.augmentation")
+            module, name = _getModule(augmentation, "konfai.data.augmentation")
             dataAugmentation: DataAugmentation = config("{}.Dataset.augmentations.{}.dataAugmentations.{}".format(KONFAI_ROOT(), key, augmentation))(getattr(importlib.import_module(module), name))(config = None)
             dataAugmentation.load(prob.prob)
             dataAugmentation.setDatasets(datasets)

konfai/data/data_manager.py

@@ -129,7 +129,7 @@ class DatasetIter(data.Dataset):
                     for group_dest in self.groups_src[group_src]:
                         self.data[group_dest][index].unloadAugmentation()
                         self.data[group_dest][index].resetAugmentation()
-        self.load(label + " Augmentation")
+            self.load(label + " Augmentation")
 
     def load(self, label: str):
         if self.use_cache:
@@ -184,7 +184,7 @@ class DatasetIter(data.Dataset):
         data = {}
         x, a, p = self.map[index]
         if x not in self._index_cache:
-            if x not in self._index_cache and len(self._index_cache) >= self.buffer_size and not self.use_cache:
+            if len(self._index_cache) >= self.buffer_size and not self.use_cache:
                 self._unloadData(self._index_cache[0])
             self._loadData(x)
 

konfai/data/patching.py

@@ -125,23 +125,20 @@ class Accumulator():
         return len(self.patch_slices) == len([v for v in self._layer_accumulator if v is not None])
 
     def assemble(self) -> torch.Tensor:
-        if all([self._layer_accumulator[0].shape[i-len(self.patch_size)] == size for i, size in enumerate(self.patch_size)]): 
-            N = 2 if self.batch else 1
-            
-            result = torch.zeros((list(self._layer_accumulator[0].shape[:N])+list(max([[v.stop for v in patch] for patch in self.patch_slices]))), dtype=self._layer_accumulator[0].dtype).to(self._layer_accumulator[0].device)
-            for patch_slice, data in zip(self.patch_slices, self._layer_accumulator):
-                slices_dest = tuple([slice(result.shape[i]) for i in range(N)] + list(patch_slice))
-
-                for dim, s in enumerate(patch_slice):
-                    if s.stop-s.start == 1:
-                        data = data.unsqueeze(dim=dim+N)
-                if self.patchCombine is not None:
-                    result[slices_dest] += self.patchCombine(data)
-                else:
-                    result[slices_dest] = data
-            result = result[tuple([slice(None, None)]+[slice(0, s) for s in self.shape])]
-        else:
-            result = torch.cat(tuple(self._layer_accumulator), dim=0)
+        N = 2 if self.batch else 1
+        result = torch.zeros((list(self._layer_accumulator[0].shape[:N])+list(max([[v.stop for v in patch] for patch in self.patch_slices]))), dtype=self._layer_accumulator[0].dtype).to(self._layer_accumulator[0].device)
+        for patch_slice, data in zip(self.patch_slices, self._layer_accumulator):
+            slices_dest = tuple([slice(result.shape[i]) for i in range(N)] + list(patch_slice))
+
+            for dim, s in enumerate(patch_slice):
+                if s.stop-s.start == 1:
+                    data = data.unsqueeze(dim=dim+N)
+            if self.patchCombine is not None:
+                result[slices_dest] += self.patchCombine(data)
+            else:
+                result[slices_dest] = data
+        result = result[tuple([slice(None, None)]+[slice(0, s) for s in self.shape])]
+
         self._layer_accumulator.clear()
         return result
 

konfai/data/transform.py

@@ -36,7 +36,7 @@ class TransformLoader:
         pass
     
     def getTransform(self, classpath : str, DL_args : str) -> Transform:
-        module, name = _getModule(classpath, "data.transform")
+        module, name = _getModule(classpath, "konfai.data.transform")
         return config("{}.{}".format(DL_args, classpath))(getattr(importlib.import_module(module), name))(config = None)
 
 class Clip(Transform):
@@ -213,7 +213,7 @@ class Resample(Transform, ABC):
 
 class ResampleToResolution(Resample):
 
-    def __init__(self, spacing : list[Union[float, None]] = [1., 1., 1.]) -> None:
+    def __init__(self, spacing : list[float] = [1., 1., 1.]) -> None:
         self.spacing = torch.tensor([0 if s < 0 else s for s in spacing])
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
@@ -244,34 +244,35 @@ class ResampleToResolution(Resample):
         cache_attribute["Size"] = np.asarray(size)
         return self._resample(input, size)
 
-class ResampleToSize(Resample):
+class ResampleToShape(Resample):
 
-    def __init__(self, size : list[int] = [100,512,512]) -> None:
-        self.size = size
+    def __init__(self, shape : list[float] = [100,256,256]) -> None:
+        self.shape = torch.tensor([0 if s < 0 else s for s in shape])
 
     def transformShape(self, shape: list[int], cache_attribute: Attribute) -> list[int]:
         if "Spacing" not in cache_attribute:
             TransformError("Missing 'Spacing' in cache attributes, the data is likely not a valid image.",
                         "Make sure your input is a image (e.g., .nii, .mha) with proper metadata.")
-        if len(shape) != len(self.size):
+        if len(shape) != len(self.shape):
             TransformError("Shape and spacing dimensions do not match: shape={shape}, spacing={self.spacing}")
-        size = self.size
-        for i, s in enumerate(self.size):
-            if s == -1:
-                size[i] = shape[i]
-        return size
+        new_shape = self.shape
+        for i, s in enumerate(self.shape):
+            if s == 0:
+                new_shape[i] = shape[i]
+        print(new_shape)
+        return new_shape
     
     def __call__(self, name: str, input: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor:
-        size = self.size
-        image_size =  np.asarray([int(x) for x in torch.tensor(input.shape[1:])])
-        for i, s in enumerate(self.size):
-            if s is None:
-                size[i] = image_size[i]
+        shape = self.shape
+        image_shape =  torch.tensor([int(x) for x in torch.tensor(input.shape[1:])])
+        for i, s in enumerate(self.shape):
+            if s == 0:
+                shape[i] = image_shape[i]
         if "Spacing" in cache_attribute:
-            cache_attribute["Spacing"] = torch.flip(torch.tensor(image_size)/torch.tensor(size)*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
-        cache_attribute["Size"] = image_size
-        cache_attribute["Size"] = size
-        return self._resample(input, size)
+            cache_attribute["Spacing"] = torch.flip(image_shape/shape*torch.flip(cache_attribute.get_tensor("Spacing"), dims=[0]), dims=[0])
+        cache_attribute["Size"] = image_shape
+        cache_attribute["Size"] = shape
+        return self._resample(input, shape)
 
 class ResampleTransform(Transform):
 

konfai/evaluator.py

@@ -27,7 +27,7 @@ class CriterionsLoader():
     def getCriterions(self, 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")
             criterions[config("{}.metrics.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
         return criterions
 

konfai/main.py

@@ -9,8 +9,6 @@ import sys
 sys.path.insert(0, os.getcwd())
 
 def main():
-    import tracemalloc
-
     parser = argparse.ArgumentParser(description="KonfAI", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     try:
         with setup(parser) as distributedObject:

konfai/metric/measure.py

@@ -3,6 +3,9 @@ import importlib
 import numpy as np
 import torch
 
+from torchvision.models import inception_v3, Inception_V3_Weights
+from torchvision.transforms import functional as F
+from scipy import linalg    
 
 import torch.nn.functional as F
 import os
@@ -247,7 +250,7 @@ class PerceptualLoss(Criterion):
         def getLoss(self) -> dict[torch.nn.Module, float]:
             result: dict[torch.nn.Module, float] = {}
             for loss, l in self.losses.items():
-                module, name = _getModule(loss, "metric.measure")
+                module, name = _getModule(loss, "konfai.metric.measure")
                 result[config(self.DL_args)(getattr(importlib.import_module(module), name))(config=None)] = l   
             return result
         
@@ -401,21 +404,42 @@ class L1LossRepresentation(Criterion):
     def forward(self, output: torch.Tensor) -> torch.Tensor:
         return self.loss(output[0], output[1])+ self._variance(output[0]) + self._variance(output[1])
 
-"""import torch
-import torch.nn as nn
-from torchvision.models import inception_v3, Inception_V3_Weights
-from torchvision.transforms import functional as F
-from scipy import linalg
-import numpy as np
+class FocalLoss(Criterion):
+
+    def __init__(self, gamma: float = 2.0, alpha: list[float] = [0.5, 2.0, 0.5, 0.5, 1], reduction: str = "mean"):
+        super().__init__()
+        raw_alpha = torch.tensor(alpha, dtype=torch.float32)
+        self.alpha = raw_alpha / raw_alpha.sum() * len(raw_alpha)
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, output: torch.Tensor, targets: list[torch.Tensor]) -> torch.Tensor:
+        target = targets[0].long()
+        
+        logpt = F.log_softmax(output, dim=1)
+        pt = torch.exp(logpt)
+
+        logpt = logpt.gather(1, target.unsqueeze(1)) 
+        pt = pt.gather(1, target.unsqueeze(1))
+
+        at = self.alpha[target].unsqueeze(1)
+        loss = -at * ((1 - pt) ** self.gamma) * logpt
+
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        return loss
     
+
 class FID(Criterion):
 
-    class InceptionV3(nn.Module):
+    class InceptionV3(torch.nn.Module):
 
         def __init__(self) -> None:
             super().__init__()
             self.model = inception_v3(weights=Inception_V3_Weights.DEFAULT, transform_input=False)
-            self.model.fc = nn.Identity()
+            self.model.fc = torch.nn.Identity()
             self.model.eval()
 
         def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -434,13 +458,11 @@ class FID(Criterion):
         return features
 
     def calculate_fid(real_features: np.ndarray, generated_features: np.ndarray) -> float:
-        # Calculate mean and covariance statistics
         mu1 = np.mean(real_features, axis=0)
         sigma1 = np.cov(real_features, rowvar=False)
         mu2 = np.mean(generated_features, axis=0)
         sigma2 = np.cov(generated_features, rowvar=False)
 
-        # Calculate FID score
         diff = mu1 - mu2
         covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
         if np.iscomplexobj(covmean):
@@ -456,9 +478,8 @@ class FID(Criterion):
         generated_features = FID.get_features(generated_images, self.inception_model)
 
         return FID.calculate_fid(real_features, generated_features)
-        """
 
-"""class MutualInformationLoss(torch.nn.Module):
+class MutualInformationLoss(torch.nn.Module):
     def __init__(self, num_bins: int = 23, sigma_ratio: float = 0.5, smooth_nr: float = 1e-7, smooth_dr: float = 1e-7) -> None:
         super().__init__()
         bin_centers = torch.linspace(0.0, 1.0, num_bins)
@@ -469,13 +490,12 @@ class FID(Criterion):
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
 
-    def parzen_windowing(self, pred: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def parzen_windowing(self, pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         pred_weight, pred_probability = self.parzen_windowing_gaussian(pred)
         target_weight, target_probability = self.parzen_windowing_gaussian(target)
         return pred_weight, pred_probability, target_weight, target_probability
 
-
-    def parzen_windowing_gaussian(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         img = torch.clamp(img, 0, 1)
         img = img.reshape(img.shape[0], -1, 1)  # (batch, num_sample, 1)
         weight = torch.exp(-self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2)  # (batch, num_sample, num_bin)
@@ -488,4 +508,4 @@ class FID(Criterion):
         pab = torch.bmm(wa.permute(0, 2, 1), wb.to(wa)).div(wa.shape[1])  # (batch, num_bins, num_bins)
         papb = torch.bmm(pa.permute(0, 2, 1), pb.to(pa))  # (batch, num_bins, num_bins)
         mi = torch.sum(pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2))  # (batch)
-        return torch.mean(mi).neg()  # average over the batch and channel ndims"""
+        return torch.mean(mi).neg()  # average over the batch and channel ndims

konfai/metric/schedulers.py

@@ -2,6 +2,7 @@ import torch
 import numpy as np
 from abc import abstractmethod
 from konfai.utils.config import config
+from functools import partial
 
 class Scheduler():
 
@@ -46,4 +47,12 @@ class CosineAnnealing(Scheduler):
         self.T_max = T_max
 
     def get_value(self):
-        return self.eta_min + (self.baseValue - self.eta_min) *(1 + np.cos(self.it * torch.pi / self.T_max)) / 2
+        return self.eta_min + (self.baseValue - self.eta_min) *(1 + np.cos(self.it * torch.pi / self.T_max)) / 2
+
+class Warmup(torch.optim.lr_scheduler.LambdaLR):
+    
+    def warmup(warmup_steps: int, step: int) -> float:
+        return min(1.0, step / warmup_steps)
+
+    def __init__(self, optimizer: torch.optim.Optimizer, warmup_steps: int = 10, last_epoch=-1, verbose="deprecated"):
+        super().__init__(optimizer, partial(Warmup.warmup, warmup_steps), last_epoch, verbose)

konfai/models/segmentation/NestedUNet.py

@@ -1,36 +1,37 @@
+from konfai.network import network, blocks
 from typing import Union
 import torch
-
-from konfai.network import network, blocks
-from konfai.utils.config import config
 from konfai.data.patching import ModelPatch
 
-     
-class NestedUNetBlock(network.ModuleArgsDict):
-
-    def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, dim: int, i : int = 0) -> None:
-        super().__init__()
-        if i > 0:
-            self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
-        
-        self.add_module("X_{}_{}".format(i, 0), block(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), out_branch=["X_{}_{}".format(i, 0)])
-        if len(channels) > 2:
-            self.add_module("UNetBlock_{}".format(i+1), NestedUNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, dim, i+1), in_branch=["X_{}_{}".format(i, 0)], out_branch=["X_{}_{}".format(i+1, j) for j in range(len(channels)-2)])
-            for j in range(len(channels)-2):
-                self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=channels[2], out_channels=channels[1], upSampleMode=upSampleMode, dim=dim), in_branch=["X_{}_{}".format(i+1, j)], out_branch=["X_{}_{}".format(i+1, j)])
-                self.add_module("SkipConnection_{}_{}".format(i, j+1), blocks.Concat(), in_branch=["X_{}_{}".format(i+1, j)]+["X_{}_{}".format(i, r) for r in range(j+1)], out_branch=["X_{}_{}".format(i, j+1)])
-                self.add_module("X_{}_{}".format(i, j+1), block(in_channels=(channels[1]*(j+1)+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*(j+2), out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), in_branch=["X_{}_{}".format(i, j+1)], out_branch=["X_{}_{}".format(i, j+1)])
-
 class NestedUNet(network.Network):
 
+    class NestedUNetBlock(network.ModuleArgsDict):
+
+        def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, block: type, dim: int, i : int = 0) -> None:
+            super().__init__()
+            if i > 0:
+                self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+            
+            self.add_module("X_{}_{}".format(i, 0), block(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), out_branch=["X_{}_{}".format(i, 0)])
+            if len(channels) > 2:
+                self.add_module("UNetBlock_{}".format(i+1), NestedUNet.NestedUNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, block, dim, i+1), in_branch=["X_{}_{}".format(i, 0)], out_branch=["X_{}_{}".format(i+1, j) for j in range(len(channels)-2)])
+                for j in range(len(channels)-2):
+                    self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=channels[2], out_channels=channels[1], upSampleMode=upSampleMode, dim=dim), in_branch=["X_{}_{}".format(i+1, j)], out_branch=["X_{}_{}".format(i+1, j)])
+                    self.add_module("SkipConnection_{}_{}".format(i, j+1), blocks.Concat(), in_branch=["X_{}_{}".format(i+1, j)]+["X_{}_{}".format(i, r) for r in range(j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+                    self.add_module("X_{}_{}".format(i, j+1), block(in_channels=(channels[1]*(j+1)+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*(j+2), out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim), in_branch=["X_{}_{}".format(i, j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+
+
     class NestedUNetHead(network.ModuleArgsDict):
 
-        def __init__(self, in_channels: int, nb_class: int, dim: int) -> None:
+        def __init__(self, in_channels: int, nb_class: int, activation: str, dim: int) -> None:
             super().__init__()
             self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
-            self.add_module("Softmax", torch.nn.Softmax(dim=1))
-            self.add_module("Argmax", blocks.ArgMax(dim=1))
-
+            if activation == "Softmax":
+                self.add_module("Softmax", torch.nn.Softmax(dim=1))
+                self.add_module("Argmax", blocks.ArgMax(dim=1))
+            elif activation == "Tanh":
+                self.add_module("Tanh", torch.nn.Tanh())
+            
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
                     schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
@@ -44,9 +45,69 @@ class NestedUNet(network.Network):
                     downSampleMode: str = "MAXPOOL",
                     upSampleMode: str = "CONV_TRANSPOSE",
                     attention : bool = False,
-                    blockType: str = "Conv") -> None:
+                    blockType: str = "Conv",
+                    activation: str = "Softmax") -> None:
         super().__init__(in_channels = channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)
                 
-        self.add_module("UNetBlock_0", NestedUNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(channels)-2)])    
+        self.add_module("UNetBlock_0", NestedUNet.NestedUNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(channels)-2)])    
         for j in range(len(channels)-2):
-            self.add_module("Head_{}".format(j), NestedUNet.NestedUNetHead(in_channels=channels[1], nb_class=nb_class, dim=dim), in_branch=["X_0_{}".format(j+1)], out_branch=[-1])
+            self.add_module("Head_{}".format(j), NestedUNet.NestedUNetHead(in_channels=channels[1], nb_class=nb_class, activation=activation, dim=dim), in_branch=["X_0_{}".format(j+1)], out_branch=[-1])
+
+
+class UNetpp(network.Network):
+
+    class ResNetEncoderLayer(network.ModuleArgsDict):
+
+        def __init__(self, in_channel: int, out_channel: int, nb_block: int, dim: int, downSampleMode : blocks.DownSampleMode):
+            super().__init__()
+            for i in range(nb_block):
+                if downSampleMode == blocks.DownSampleMode.MAXPOOL and i == 0:
+                    self.add_module("DownSample", blocks.getTorchModule("MaxPool", dim)(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False))
+                self.add_module("ResBlock_{}".format(i), blocks.ResBlock(in_channel, out_channel, [blocks.BlockConfig(3, 2 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i == 0 else 1, 1, False, "ReLU;True", blocks.NormMode.BATCH), blocks.BlockConfig(3, 1, 1, False, None, blocks.NormMode.BATCH)], dim=dim))
+                in_channel = out_channel
+            
+    def resNetEncoder(channels: list[int], layers: list[int], dim: int) -> list[torch.nn.Module]:
+        modules = []
+        modules.append(blocks.ConvBlock(channels[0], channels[1], [blocks.BlockConfig(7, 2, 3, False, "ReLU", blocks.NormMode.BATCH)], dim=dim))
+        for i, (in_channel, out_channel, layer) in enumerate(zip(channels[1:], channels[2:], layers)):
+            modules.append(UNetpp.ResNetEncoderLayer(in_channel, out_channel, layer, dim, blocks.DownSampleMode.MAXPOOL if i==0 else blocks.DownSampleMode.CONV_STRIDE))
+        return modules
+
+    class UNetPPBlock(network.ModuleArgsDict):
+
+        def __init__(self, encoder_channels: list[int], decoder_channels: list[int], encoders: list[torch.nn.Module], upSampleMode: blocks.UpSampleMode, dim: int, i : int = 0) -> None:
+            super().__init__()
+            self.add_module("X_{}_{}".format(i, 0), encoders[0], out_branch=["X_{}_{}".format(i, 0)])
+            if len(encoder_channels) > 2:
+                self.add_module("UNetBlock_{}".format(i+1), UNetpp.UNetPPBlock(encoder_channels[1:], decoder_channels[1:], encoders[1:], upSampleMode, dim, i+1), in_branch=["X_{}_{}".format(i, 0)], out_branch=["X_{}_{}".format(i+1, j) for j in range(len(encoder_channels)-2)])
+                for j in range(len(encoder_channels)-2):
+                    self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=encoder_channels[2], out_channels=decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1], upSampleMode=upSampleMode, dim=dim), in_branch=["X_{}_{}".format(i+1, j)], out_branch=["X_{}_{}".format(i+1, j)])
+                    self.add_module("SkipConnection_{}_{}".format(i, j+1), blocks.Concat(), in_branch=["X_{}_{}".format(i+1, j)]+["X_{}_{}".format(i, r) for r in range(j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+                    self.add_module("X_{}_{}".format(i, j+1), blocks.ConvBlock(in_channels=encoder_channels[1]*j+((encoder_channels[1]+encoder_channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else (decoder_channels[1]+decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1]+encoder_channels[1])), out_channels=decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1], blockConfigs=[blocks.BlockConfig(3,1,1,False, "ReLU;True", blocks.NormMode.BATCH)]*2, dim=dim), in_branch=["X_{}_{}".format(i, j+1)], out_branch=["X_{}_{}".format(i, j+1)])
+
+    class UNetPPHead(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, nb_class: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Upsample", blocks.upSample(in_channels=in_channels, out_channels=out_channels, upSampleMode=blocks.UpSampleMode.UPSAMPLE_BILINEAR, dim=dim))
+            self.add_module("ConvBlock", blocks.ConvBlock(in_channels=in_channels, out_channels=out_channels, blockConfigs=[blocks.BlockConfig(3,1,1,False, "ReLU;True", blocks.NormMode.BATCH)]*2, dim=dim))
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = out_channels, out_channels = nb_class, kernel_size = 3, stride = 1, padding = 1))
+            if nb_class > 1:
+                self.add_module("Softmax", torch.nn.Softmax(dim=1))
+                self.add_module("Argmax", blocks.ArgMax(dim=1))
+            else:
+                self.add_module("Tanh", torch.nn.Tanh())
+                
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    encoder_channels: list[int] = [1,64,64,128,256,512],
+                    decoder_channels: list[int] = [256,128,64,32,16,1],
+                    layers: list[int] = [3,4,6,3],
+                    dim : int = 2) -> None:
+        super().__init__(in_channels = encoder_channels[0], optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, patch=patch, dim = dim)
+        self.add_module("Block_0", UNetpp.UNetPPBlock(encoder_channels, decoder_channels[::-1], UNetpp.resNetEncoder(encoder_channels, layers, dim), blocks.UpSampleMode.UPSAMPLE_BILINEAR, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(encoder_channels)-2)])    
+        self.add_module("Head", UNetpp.UNetPPHead(in_channels=decoder_channels[-3], out_channels=decoder_channels[-2], nb_class=decoder_channels[-1], dim=dim), in_branch=["X_0_{}".format(len(encoder_channels)-2)], out_branch=[-1])
+

konfai/models/segmentation/UNet.py

@@ -12,7 +12,6 @@ class UNetHead(network.ModuleArgsDict):
         self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = nb_class, kernel_size = 1, stride = 1, padding = 0))
         self.add_module("Softmax", torch.nn.Softmax(dim=1))
         self.add_module("Argmax", blocks.ArgMax(dim=1))
-        
 
 class UNetBlock(network.ModuleArgsDict):
 

konfai/network/blocks.py

@@ -19,7 +19,7 @@ class NormMode(Enum):
     SYNCBATCH = 5
     INSTANCE_AFFINE = 6
 
-def getNorm(normMode: Enum, channels : int, dim: int) -> torch.nn.Module:
+def getNorm(normMode: Enum, channels : int, dim: int) -> Union[torch.nn.Module, None]:
     if normMode == NormMode.BATCH:
         return getTorchModule("BatchNorm", dim = dim)(channels, affine=True, track_running_stats=True)
     if normMode == NormMode.INSTANCE:
@@ -32,7 +32,7 @@ def getNorm(normMode: Enum, channels : int, dim: int) -> torch.nn.Module:
         return torch.nn.GroupNorm(num_groups=32, num_channels=channels)
     if normMode == NormMode.LAYER:
         return torch.nn.GroupNorm(num_groups=1, num_channels=channels)
-    return torch.nn.Identity()
+    return None
 
 class UpSampleMode(Enum):
     CONV_TRANSPOSE = 0,
@@ -59,9 +59,9 @@ class BlockConfig():
         self.stride = stride
         self.padding = padding
         self.activation = activation
-        self.normMode = normMode
-
-        if isinstance(normMode, str):
+        if normMode is None:
+            self.norm = None
+        elif isinstance(normMode, str):
             self.norm = NormMode._member_map_[normMode]
         elif isinstance(normMode, NormMode):
             self.norm = normMode
@@ -70,9 +70,13 @@ class BlockConfig():
         return getTorchModule("Conv", dim = dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = self.kernel_size, stride = self.stride, padding = self.padding, bias=self.bias)
     
     def getNorm(self, channels : int, dim: int) -> torch.nn.Module:
+        if self.norm is None:
+            return None
         return getNorm(self.norm, channels, dim) if isinstance(self.norm, NormMode) else self.norm(channels)
 
     def getActivation(self) -> torch.nn.Module:
+        if self.activation is None:
+            return None
         if isinstance(self.activation, str):
             return getTorchModule(self.activation.split(";")[0])(*[ast.literal_eval(value) for value in self.activation.split(";")[1:]]) if self.activation != "None" else torch.nn.Identity()
         return self.activation()
@@ -83,21 +87,35 @@ class ConvBlock(network.ModuleArgsDict):
         super().__init__()
         for i, blockConfig in enumerate(blockConfigs):
             self.add_module("Conv_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[0])
-            self.add_module("Norm_{}".format(i), blockConfig.getNorm(out_channels, dim), alias=alias[1])
-            self.add_module("Activation_{}".format(i), blockConfig.getActivation(), alias=alias[2])
+            norm = blockConfig.getNorm(out_channels, dim)
+            if norm is not None:
+                self.add_module("Norm_{}".format(i), norm, alias=alias[1])
+            activation = blockConfig.getActivation()
+            if activation is not None:
+                self.add_module("Activation_{}".format(i), activation, alias=alias[2])
             in_channels = out_channels
 
 class ResBlock(network.ModuleArgsDict):
     
-    def __init__(self, in_channels : int, out_channels : int, nb_conv: int, blockConfig : BlockConfig, dim : int, alias : list[list[str]]=[[], [], [], []]) -> None:
+    def __init__(self, in_channels : int, out_channels : int, blockConfigs : list[BlockConfig], dim : int, alias : list[list[str]]=[[], [], [], [], []]) -> None:
         super().__init__()
-        for i in range(nb_conv):
-            self.add_module("ConvBlock_{}".format(i), ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blockConfig, dim=dim, alias=alias[:3]))
+        for i, blockConfig in enumerate(blockConfigs):
+            self.add_module("Conv_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[0])
+            norm = blockConfig.getNorm(out_channels, dim)
+            if norm is not None:
+                self.add_module("Norm_{}".format(i), norm, alias=alias[1])
+            activation = blockConfig.getActivation()
+            if activation is not None:
+                self.add_module("Activation_{}".format(i), activation, alias=alias[2])
+
             if in_channels != out_channels:
-                self.add_module("Conv_skip_{}".format(i), blockConfig.getConv(in_channels, out_channels, dim), alias=alias[3], in_branch=[1], out_branch=[1])    
+                self.add_module("Conv_skip", getTorchModule("Conv", dim)(in_channels, out_channels, 1, blockConfig.stride, bias=blockConfig.bias), alias=alias[3], in_branch=[1], out_branch=[1])    
+                self.add_module("Norm_skip", blockConfig.getNorm(out_channels, dim), alias=alias[4], in_branch=[1], out_branch=[1])    
             in_channels = out_channels
-            self.add_module("Add_{}".format(i), Add(), in_branch=[0,1], out_branch=[0,1])
-    
+            
+        self.add_module("Add", Add(), in_branch=[0,1])
+        self.add_module("Norm_{}".format(i+1), torch.nn.ReLU(inplace=True))
+            
 def downSample(in_channels: int, out_channels: int, downSampleMode: DownSampleMode, dim: int) -> torch.nn.Module:
     if downSampleMode == DownSampleMode.MAXPOOL:
         return getTorchModule("MaxPool", dim = dim)(2)
@@ -176,6 +194,7 @@ 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, MeasureError
+from konfai.utils.utils import State, _getModule, getDevice, getGPUMemory, MeasureError, TrainerError
 from konfai.data.patching import Accumulator, ModelPatch
 
 class NetState(Enum):
@@ -55,9 +55,16 @@ class LRSchedulersLoader():
 
     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
+        for nameTmp, step in self.params.items():
+            if nameTmp:
+                ok = False
+                for m in ["torch.optim.lr_scheduler", "konfai.metric.schedulers"]:
+                    module, name = _getModule(nameTmp, m)
+                    if hasattr(importlib.import_module(module), name):
+                        schedulers[config("{}.Model.{}.Schedulers.{}".format(KONFAI_ROOT(), key, name))(getattr(importlib.import_module(module), name))(optimizer, config = None)] = step.nb_step        
+                        ok = True
+                if not ok:
+                    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(nameTmp))
         return schedulers
 
 class SchedulersLoader():
@@ -95,7 +102,7 @@ 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))
             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
@@ -133,9 +140,10 @@ class Measure():
             self._loss.clear()
 
         def add(self, weight: float, value: torch.Tensor) -> None:
-            self._loss.append(value if self.isLoss else value.detach())
-            self._values.append(value.item())
-            self._weight.append(weight)
+            if self.isLoss or value != 0:
+                self._loss.append(value if self.isLoss else value.detach())
+                self._values.append(value.item())
+                self._weight.append(weight)
 
         def getLastLoss(self) -> torch.Tensor:
             return self._loss[-1]*self._weight[-1] if len(self._loss) else torch.zeros((1), requires_grad = True)
@@ -560,6 +568,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()
@@ -654,9 +663,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())])
+        if self.schedulers:
+            for scheduler in self.schedulers:
+                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]:
@@ -885,11 +900,12 @@ 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):
+                if value is None or (self._nb_lr_update >= step and self._nb_lr_update < step+value):
                     break
                 step += value
         if scheduler:
@@ -897,7 +913,7 @@ class Network(ModuleArgsDict, ABC):
                 if self.measure:
                     scheduler.step(sum(self.measure.getLastValues(0).values()))
             else:
-                scheduler.step()     
+                scheduler.step()
 
     @_function_network()
     def getNetworks(self) -> Self:
@@ -915,6 +931,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:
@@ -943,7 +965,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

@@ -8,7 +8,7 @@ import os
 
 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.utils import State, get_patch_slices_from_nb_patch_per_dim, NeedDevice, _getModule, DistributedObject, DataLog, description, PredictorError
 from konfai.utils.dataset import Dataset, Attribute
 from konfai.data.data_manager import DataPrediction, DatasetIter
 from konfai.data.patching import Accumulator, PathCombine
@@ -46,7 +46,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
         self.names: dict[int, str] = {}
         self.nb_data_augmentation = 0
 
-    def load(self, name_layer: str, datasets: list[Dataset]):
+    def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
         transforms_type = ["pre_transforms", "post_transforms", "final_transforms"]
         for name, _transform_type, transform_type in [(k, getattr(self, "_{}".format(k)), getattr(self, k)) for k in transforms_type]:
             
@@ -57,7 +57,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
                     transform_type.append(transform)
 
         if self._patchCombine is not None:
-            module, name = _getModule(self._patchCombine, "data.patching")
+            module, name = _getModule(self._patchCombine, "konfai.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:
@@ -96,17 +96,21 @@ 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):
 
@@ -114,7 +118,7 @@ class OutSameAsGroupDataset(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)
         self.group_src, self.group_dest = sameAsGroup.split(":")
-        self.reduction = reduction
+        self.reduction_classpath = reduction
         self.inverse_transform = inverse_transform
 
     def addLayer(self, index_dataset: int, index_augmentation: int, index_patch: int, layer: torch.Tensor, dataset: DatasetIter):
@@ -137,9 +141,24 @@ class OutSameAsGroupDataset(OutDataset):
         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])
-                
         self.output_layer_accumulator[index_dataset][index_augmentation].addLayer(index_patch, layer)
 
+
+    def load(self, name_layer: str, datasets: list[Dataset], groups: dict[str, str]):
+        super().load(name_layer, datasets, groups)
+        module, name = _getModule(self.reduction_classpath, "konfai.predictor")
+        self.reduction = config("{}.outsDataset.{}.OutDataset.{}".format(KONFAI_ROOT(), name_layer, self.reduction_classpath))(getattr(importlib.import_module(module), name))(config = None)
+       
+        if self.group_src not in groups.keys():
+            raise PredictorError(
+                f"Source group '{self.group_src}' not found. Available groups: {list(groups.keys())}."
+            )
+
+        if self.group_dest not in groups[self.group_src]:
+            raise PredictorError(
+                f"Destination group '{self.group_dest}' not found. Available groups: {groups[self.group_src]}."
+            )
+    
     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]
@@ -166,14 +185,8 @@ class OutSameAsGroupDataset(OutDataset):
         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)
 
-        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)")
         for transform in self.final_transforms:
             result = transform(name, result, self.attributes[index][0][0])
         return result
@@ -294,9 +307,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)])
 
@@ -317,12 +330,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
     
@@ -345,7 +353,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)
@@ -403,10 +411,24 @@ class Predictor(DistributedObject):
 
         shutil.copyfile(CONFIG_FILE(), self.predict_path+"Prediction.yml")
 
+        
         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)
+        
+        modules = []
+        for i,_,_ in self.model.named_ModuleArgsDict():
+            modules.append(i)
+        for output_group in self.outsDataset.keys():
+            if output_group not in modules:
+                raise PredictorError("The output group '{}' defined in 'outputsCriterions' does not correspond to any module in the model.".format(output_group),
+                    "Available modules: {}".format(modules),
+                    "Please check that the name matches exactly a submodule or output of your model architecture."
+                )
+        module, name = _getModule(self.combine_classpath, "konfai.predictor")
+        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:
@@ -415,7 +437,7 @@ class Predictor(DistributedObject):
         self.size = (len(self.gpu_checkpoints)+1 if self.gpu_checkpoints else 1)
         self.dataloader = self.dataset.getData(world_size//self.size)
         for name, outDataset in self.outsDataset.items():
-            outDataset.load(name.replace(".", ":"), list(self.dataset.datasets.values()))
+            outDataset.load(name.replace(".", ":"), list(self.dataset.datasets.values()), {src : dest for src, inner in self.dataset.groups_src.items() for dest in inner})
            
            
     def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):

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
     
@@ -120,7 +118,8 @@ class _Trainer():
 
     def run(self) -> None:
         self.dataloader_training.dataset.load("Train")
-        self.dataloader_validation.dataset.load("Validation")
+        if self.dataloader_validation is not None:
+            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.train()
@@ -211,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":
@@ -355,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

@@ -37,7 +37,7 @@ def _getModule(classpath : str, type : str) -> tuple[str, str]:
         module = ".".join(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
 
@@ -586,6 +586,11 @@ class EvaluatorError(KonfAIError):
     def __init__(self, *message) -> None:
         super().__init__("Evaluator", message)
 
+class PredictorError(KonfAIError):
+
+    def __init__(self, *message) -> None:
+        super().__init__("Predictor", message)
+
 class TransformError(KonfAIError):
 
     def __init__(self, *message) -> None:

{konfai-1.1.4.dist-info → konfai-1.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.4
+Version: 1.1.6
 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.4.dist-info → konfai-1.1.6.dist-info}/RECORD

@@ -1,16 +1,16 @@
 konfai/__init__.py,sha256=YXG-wpSEXWs6Jt3BDI77V4r89gEUNX-6lxW9btj5VYI,851
-konfai/evaluator.py,sha256=G8tXI-g8yTxBZXmESpRW7B5UEI4SMXq47fsGksU0G0M,8394
-konfai/main.py,sha256=xhCIs3VKXxNCgyyN616K-VANEFCsZZRR-l8LIMHlmPw,2597
-konfai/predictor.py,sha256=R9ohOqEcSoq6Wu9Yv3pkerVfm6qSotd4aICe7ctOsU0,22728
-konfai/trainer.py,sha256=gJ7qmU302M3Mevq3l1aIdPpey5C3prtibD69LyxbnAg,20240
+konfai/evaluator.py,sha256=WM78NGydV0iqKElahr-WpOCZLEJXHjjPq-LS-gi23rk,8401
+konfai/main.py,sha256=kr7Iie_f67NF6G3dAAj9G6Z9dhn9RzbdLYpzy2WvIh8,2573
+konfai/predictor.py,sha256=kqAF3P4BAOZvjHtBAfv2M4-yonqAOcnIV7APgxdRsww,24283
+konfai/trainer.py,sha256=Edi8l8OOfCewYM-Cd5C5rCqCaprvlfxzohd4iLkK5u0,20632
 konfai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/data/augmentation.py,sha256=ruvS7ykFveOZBe2Mv_TOU5a8AzUgH7gfTipwqYgRM7k,32293
-konfai/data/data_manager.py,sha256=TCQFZ_zpRHYIomPOotG8U_I0DfpjhSaajb7ypd6fyRU,30103
-konfai/data/patching.py,sha256=oq43r3JlS9hw-pnzpjT3epvXJxEKqiQEtFVZPjY9Jic,15749
-konfai/data/transform.py,sha256=ZZoxft0V8t-DaKjF3OsIFezz41B8_xwPqBvYCYwF0dk,26471
+konfai/data/augmentation.py,sha256=SD8m0uLz5aQYtjAjns0H04WgrDWSN6nd7lWlNnT2fi8,32300
+konfai/data/data_manager.py,sha256=eYgFJjHr9HEpRLTtLabUZLPH4c82OQCiGuYQ84ktWCE,30076
+konfai/data/patching.py,sha256=zs3T4yTV8_iCFrqO21bo6GhwTywoTxIL31IAi-jiJDQ,15478
+konfai/data/transform.py,sha256=LfdsNUE5_6mIGrVzI6B2N2WO-jaNzFmuMxhP1Y0Fbwk,26538
 konfai/metric/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/metric/measure.py,sha256=K7s15wNJm1_Iaav9y7Oe_UhSnbLUiybhLTVumHZ8ZaY,22001
-konfai/metric/schedulers.py,sha256=UoSr1TW_hrus3DhvOEbDefxCSUGz7lJS_8vbz0GEye8,1370
+konfai/metric/measure.py,sha256=j-rEsNgVMPjYFWtahYEaiMSDnyD5Dv8edx5KEt2Rj2U,22797
+konfai/metric/schedulers.py,sha256=eZIlJMOQH8AV48F5hiRFbmONmoY4hHMujvS3qgFYXbw,1771
 konfai/models/classification/convNeXt.py,sha256=Phj1hO8TCItVhBXoFXQcIkA85DZxurGLyHIPWBfXJ0Y,9243
 konfai/models/classification/resnet.py,sha256=t8KJEgGudWBpGJM9YS1lH_5eX0-wGbK7ZW5uZW214eM,7958
 konfai/models/generation/cStyleGan.py,sha256=PwaVQX9h5u_8YbbxA7sz7inzIEoMadinuo0kFFNqIn0,8042
@@ -20,20 +20,20 @@ konfai/models/generation/gan.py,sha256=-GoKxHm3W9NdD4U77UcJrG5TfOZ3NWFUZG663kt2X
 konfai/models/generation/vae.py,sha256=_3JYVT2ojZ0P98tYcD2ny7a-gWVUmnByLDhY7i-n_4g,4719
 konfai/models/registration/registration.py,sha256=18EiWt4RJIXLyFtqU-kHjV1sMnQRm9mxAA6_-2B1YqI,6313
 konfai/models/representation/representation.py,sha256=RwQYoxtdph440-t_ZLelykl0hkUAD1zdspQaLkgxb-0,2677
-konfai/models/segmentation/NestedUNet.py,sha256=wutcHYpZkOm8FqPPt5zifD6dGSDlhB9zIBlibzVyAYY,4275
-konfai/models/segmentation/UNet.py,sha256=Uei51-_0Le24ZJs9tbuL6hQn66kigkfbaLSMlH38els,4045
+konfai/models/segmentation/NestedUNet.py,sha256=D2qChm5SSYqV7FayCe9eT260vxckxRfBOCS1US6xcnw,10184
+konfai/models/segmentation/UNet.py,sha256=TKPmhHEnlOYfnUF2Qof3VSETl5nEIh9h0SkblBRdbbg,4036
 konfai/network/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-konfai/network/blocks.py,sha256=RaTI0Lrvq1V-GIFei-WTUB6wlg4LydZksAyJ8DMk40M,13502
-konfai/network/network.py,sha256=YCPfgQ7EuRiTgAlMlQTx5ospluO4aiqSbsdU4GdLiMA,47467
+konfai/network/blocks.py,sha256=P7mEuuE1B0HRSgXRCeLxZvzYGKLskHD33hskpdkjIPs,14370
+konfai/network/network.py,sha256=zszV7z6tyZpZtWb9SHXqH1pJlvdy9YB4yrf7Muv_HiM,48936
 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=-W5fbsYsTh4wIOEZBEQY7krZvGU7JC8tXHYC2rZAkaM,23459
-konfai-1.1.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-konfai-1.1.4.dist-info/METADATA,sha256=86Azx-ool5VrcYT48km-OOMvvBQzjJX1jZihzDxAhgs,2515
-konfai-1.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-konfai-1.1.4.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
-konfai-1.1.4.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
-konfai-1.1.4.dist-info/RECORD,,
+konfai/utils/utils.py,sha256=1M46MYSZuPrbH5ihu8WZHaQNfqmvfeVsH-mAk0ED5TI,23575
+konfai-1.1.6.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+konfai-1.1.6.dist-info/METADATA,sha256=wU8CbR5KO85P9dY3asD_ylG5rMCFIAQmNL0ifTDXA8g,2515
+konfai-1.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+konfai-1.1.6.dist-info/entry_points.txt,sha256=fG82HRN5-g39ACSOCtij_I3N6EHxfYnMR0D7TI_8pW8,81
+konfai-1.1.6.dist-info/top_level.txt,sha256=xF470dkIlFoFqTZEOlRehKJr4WU_8OKGXrJqYm9vWKs,7
+konfai-1.1.6.dist-info/RECORD,,