konfai 1.1.4__tar.gz → 1.1.5__tar.gz

{konfai-1.1.4 → konfai-1.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.4
+Version: 1.1.5
 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 → konfai-1.1.5}/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:

{konfai-1.1.4 → konfai-1.1.5}/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-1.1.4 → konfai-1.1.5}/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
@@ -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:
@@ -456,9 +480,9 @@ 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 +493,13 @@ 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 +512,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-1.1.5/konfai/models/segmentation/NestedUNet.py

@@ -0,0 +1,115 @@
+from konfai.network import network, blocks
+from typing import Union
+import torch
+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 NestedUNetHead(network.ModuleArgsDict):
+
+        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))
+            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(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    dim : int = 3,
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    nb_class: int = 2,
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    attention : bool = False,
+                    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)])    
+        for j in range(len(channels)-2):
+            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 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(ResNetEncoderLayer(in_channel, out_channel, layer, dim, blocks.DownSampleMode.MAXPOOL if i==0 else blocks.DownSampleMode.CONV_STRIDE))
+    return modules
+
+class NestedUNetBlock(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), NestedUNetBlock(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 NestedUNetHead(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())
+        
+class UNetpp(network.Network):
+
+
+    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", NestedUNetBlock(encoder_channels, decoder_channels[::-1], 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", NestedUNetHead(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-1.1.4 → konfai-1.1.5}/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-1.1.4 → konfai-1.1.5}/konfai/network/network.py

@@ -133,9 +133,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)

{konfai-1.1.4 → konfai-1.1.5}/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]:
             
@@ -137,9 +137,22 @@ 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)
+
+        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]
@@ -403,9 +416,20 @@ 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)
+        
+        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."
+                )
         self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
         self.modelComposite.load(self._load())
 
@@ -415,7 +439,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-1.1.4 → konfai-1.1.5}/konfai/trainer.py

@@ -120,7 +120,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()

{konfai-1.1.4 → konfai-1.1.5}/konfai/utils/utils.py

@@ -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 → konfai-1.1.5}/konfai.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: konfai
-Version: 1.1.4
+Version: 1.1.5
 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 → konfai-1.1.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "konfai"
-version = "1.1.4"
+version = "1.1.5"
 description = "Modular and configurable Deep Learning framework with YAML and PyTorch"
 readme = "README.md"
 requires-python = ">=3.8"

konfai-1.1.4/konfai/models/segmentation/NestedUNet.py

@@ -1,52 +0,0 @@
-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 NestedUNetHead(network.ModuleArgsDict):
-
-        def __init__(self, in_channels: int, nb_class: int, 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))
-
-    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,
-                    dim : int = 3,
-                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
-                    nb_class: int = 2,
-                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
-                    nb_conv_per_stage: int = 2,
-                    downSampleMode: str = "MAXPOOL",
-                    upSampleMode: str = "CONV_TRANSPOSE",
-                    attention : bool = False,
-                    blockType: str = "Conv") -> 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)])    
-        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])