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

konfai/models/generation/diffusionGan.py

@@ -37,10 +37,9 @@ class Discriminator(network.Network):
             self.add_module("Layers", Discriminator.DiscriminatorNLayers(channels, strides, dim))
             self.add_module("Head", Discriminator.DiscriminatorHead(channels[-1], dim))
 
-    @config("Discriminator")
     def __init__(self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     channels: list[int] = [1, 16, 32, 64, 64],
                     strides: list[int] = [2,2,2,1],
@@ -169,10 +168,9 @@ class Discriminator_ADA(network.Network):
             self.add_module("Layers", Discriminator_ADA.DiscriminatorNLayers(channels, strides, 100, dim), in_branch=[0, "te"])
             self.add_module("Head", Discriminator_ADA.DiscriminatorHead(channels[-1], dim))
 
-    @config("Discriminator_ADA")
     def __init__(self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     channels: list[int] = [1, 16, 32, 64, 64],
                     strides: list[int] = [2,2,2,1],
@@ -261,10 +259,9 @@ class Discriminator_ADA(network.Network):
             self.add_module("AutoEncoder", GeneratorV1.GeneratorAutoEncoder(ngf, dim))
             self.add_module("Head", GeneratorV1.GeneratorHead(in_channels=ngf, out_channels=1, dim=dim))
 
-    @config("GeneratorV1")
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     patch : ModelPatch = ModelPatch(),
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     dim : int = 3) -> None:
@@ -349,10 +346,9 @@ class GeneratorV1(network.Network):
             self.add_module("AutoEncoder", GeneratorV1.GeneratorAutoEncoder(ngf, dim))
             self.add_module("Head", GeneratorV1.GeneratorHead(in_channels=ngf, out_channels=1, dim=dim))
 
-    @config("GeneratorV1")
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     patch : ModelPatch = ModelPatch(),
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     dim : int = 3) -> None:
@@ -383,10 +379,9 @@ class GeneratorV2(network.Network):
             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)])    
             self.add_module("Head", GeneratorV2.NestedUNetHead(channels[:2], dim=dim), in_branch=["X_0_{}".format(len(channels)-2)])
 
-    @config("GeneratorV2")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     channels: list[int]=[1, 64, 128, 256, 512, 1024],
@@ -424,10 +419,9 @@ class GeneratorV3(network.Network):
             self.add_module("UNetBlock_0", UNet.UNetBlock(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, nb_class=1, dim=dim), out_branch=["X_0_{}".format(j+1) for j in range(len(channels)-2)])    
             self.add_module("Head", GeneratorV3.NestedUNetHead(channels[:2], dim=dim), in_branch=["X_0_{}".format(len(channels)-2)])
 
-    @config("GeneratorV3")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     channels: list[int]=[1, 64, 128, 256, 512, 1024],
@@ -443,7 +437,6 @@ class GeneratorV3(network.Network):
 
 class DiffusionGan(network.Network):
 
-    @config("DiffusionGan")
     def __init__(self, generator : GeneratorV1 = GeneratorV1(), discriminator : Discriminator_ADA = Discriminator_ADA()) -> None:
         super().__init__()
         self.add_module("Generator_A_to_B", generator, in_branch=[0], out_branch=["pB"])
@@ -454,7 +447,6 @@ class DiffusionGan(network.Network):
 
 class DiffusionGanV2(network.Network):
 
-    @config("DiffusionGan")
     def __init__(self, generator : GeneratorV2 = GeneratorV2(), discriminator : Discriminator = Discriminator()) -> None:
         super().__init__()
         self.add_module("Generator_A_to_B", generator, in_branch=[0], out_branch=["pB"])
@@ -466,10 +458,9 @@ class DiffusionGanV2(network.Network):
 
 class CycleGanDiscriminator(network.Network):
 
-    @config("CycleGanDiscriminator")
     def __init__(self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     channels: list[int] = [1, 16, 32, 64, 64],
@@ -485,10 +476,9 @@ class CycleGanDiscriminator(network.Network):
 
 class CycleGanGeneratorV1(network.Network):
 
-    @config("CycleGanGeneratorV1")
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     dim : int = 3) -> None:
@@ -498,10 +488,9 @@ class CycleGanGeneratorV1(network.Network):
 
 class CycleGanGeneratorV2(network.Network):
 
-    @config("CycleGanGeneratorV2")
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     channels: list[int]=[1, 64, 128, 256, 512, 1024],
@@ -518,10 +507,9 @@ class CycleGanGeneratorV2(network.Network):
 
 class CycleGanGeneratorV3(network.Network):
 
-    @config("CycleGanGeneratorV3")
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     channels: list[int]=[1, 64, 128, 256, 512, 1024],
@@ -538,7 +526,6 @@ class CycleGanGeneratorV3(network.Network):
 
 class DiffusionCycleGan(network.Network):
 
-    @config("DiffusionCycleGan")
     def __init__(self, generators : CycleGanGeneratorV3 = CycleGanGeneratorV3(), discriminators : CycleGanDiscriminator = CycleGanDiscriminator()) -> None:
         super().__init__()
         self.add_module("Generator", generators, in_branch=[0, 1], out_branch=["pB", "pA"])

konfai/models/generation/gan.py

@@ -23,10 +23,9 @@ class Discriminator(network.Network):
             self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(tuple([1]*dim)))
             self.add_module("Flatten", torch.nn.Flatten(1))
     
-    @config("Discriminator")
     def __init__(self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     nb_batch_per_step: int = 64,
                     dim : int = 3) -> None:
@@ -101,11 +100,10 @@ class Generator(network.Network):
             self.add_module("Encoder", Generator.GeneratorEncoder(channels, dim))
             self.add_module("NResBlock", Generator.GeneratorNResnetBlock(channels=channels[-1], nb_conv=6, dim=dim))
             self.add_module("Decoder", Generator.GeneratorDecoder(channels, dim))
-            
-    @config("Generator")
+        
     def __init__(self, 
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     patch : ModelPatch = ModelPatch(),
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     nb_batch_per_step: int = 64,
@@ -121,7 +119,6 @@ class Generator(network.Network):
 
 class Gan(network.Network):
 
-    @config("Gan")
     def __init__(self, generator : Generator = Generator(), discriminator : Discriminator = Discriminator()) -> None:
         super().__init__()
         self.add_module("Discriminator_B", discriminator, in_branch=[1], out_branch=[-1], requires_grad=True)

konfai/models/generation/vae.py

@@ -25,7 +25,6 @@ class VAE(network.Network):
             self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = 3, stride = 1, padding = 1))
             self.add_module("Tanh", torch.nn.Tanh()) 
 
-    @config("VAE")
     def __init__(self,
                     optimizer: network.OptimizerLoader = network.OptimizerLoader(),
                     schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),
@@ -60,7 +59,6 @@ class LinearVAE(network.Network):
             self.add_module("Linear", torch.nn.Linear(in_features, out_features))
             self.add_module("Tanh", torch.nn.Tanh())
 
-    @config("LinearVAE")
     def __init__(self,
                     optimizer: network.OptimizerLoader = network.OptimizerLoader(),
                     schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),

konfai/models/registration/registration.py

@@ -8,10 +8,9 @@ from konfai.models.segmentation import UNet
 
 class VoxelMorph(network.Network):
                 
-    @config("VoxelMorph")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     dim : int = 3,
                     channels : list[int] = [4, 16,32,32,32],

konfai/models/representation/representation.py

@@ -47,10 +47,9 @@ class Adaptation(torch.nn.Module):
 
 class Representation(network.Network):
 
-    @config("Representation")
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     dim : int = 3):
         super().__init__(in_channels = 1, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, init_type="kaiming")

konfai/models/segmentation/NestedUNet.py

@@ -34,7 +34,7 @@ class NestedUNet(network.Network):
             
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     dim : int = 3,
@@ -81,15 +81,17 @@ class UNetpp(network.Network):
             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)])
+                    in_channels = decoder_channels[3] if j == len(encoder_channels)-3 and len(encoder_channels) > 3 else encoder_channels[2]
+                    out_channel = decoder_channels[2] if j == len(encoder_channels)-3 else encoder_channels[1]
+                    self.add_module("X_{}_{}_{}".format(i, j+1, upSampleMode.name), blocks.upSample(in_channels=in_channels, out_channels=out_channel, 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)])
-
+                    self.add_module("X_{}_{}".format(i, j+1), blocks.ConvBlock(in_channels=encoder_channels[1]*(j+1)+(in_channels if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else out_channel), out_channels=out_channel, 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("Upsample", blocks.upSample(in_channels=in_channels, out_channels=out_channels, upSampleMode=blocks.UpSampleMode.UPSAMPLE, 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:
@@ -100,7 +102,7 @@ class UNetpp(network.Network):
                 
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)}, 
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     encoder_channels: list[int] = [1,64,64,128,256,512],
@@ -108,6 +110,5 @@ class UNetpp(network.Network):
                     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])
-
+        self.add_module("Block_0", UNetpp.UNetPPBlock(encoder_channels, decoder_channels[::-1], UNetpp.resNetEncoder(encoder_channels, layers, dim), blocks.UpSampleMode.CONV_TRANSPOSE, 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

@@ -39,7 +39,7 @@ class UNet(network.Network):
 
     def __init__(   self,
                     optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
                     outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
                     patch : Union[ModelPatch, None] = None,
                     dim : int = 3,

konfai/network/blocks.py

@@ -36,11 +36,7 @@ def getNorm(normMode: Enum, channels : int, dim: int) -> Union[torch.nn.Module,
 
 class UpSampleMode(Enum):
     CONV_TRANSPOSE = 0,
-    UPSAMPLE_NEAREST = 1,
-    UPSAMPLE_LINEAR = 2,
-    UPSAMPLE_BILINEAR = 3,
-    UPSAMPLE_BICUBIC = 4,
-    UPSAMPLE_TRILINEAR = 5
+    UPSAMPLE = 1,
 
 class DownSampleMode(Enum):
     MAXPOOL = 0,
@@ -128,7 +124,13 @@ def upSample(in_channels: int, out_channels: int, upSampleMode: UpSampleMode, di
     if upSampleMode == UpSampleMode.CONV_TRANSPOSE:
         return getTorchModule("ConvTranspose", dim = dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = kernel_size, stride = stride, padding = 0)
     else:
-        return torch.nn.Upsample(scale_factor=2, mode=upSampleMode.name.replace("UPSAMPLE_", "").lower(), align_corners=False)
+        if dim == 3:
+            upsampleMethod = "trilinear" 
+        if dim == 2:
+            upsampleMethod = "bilinear"
+        if dim == 1:
+            upsampleMethod = "linear"
+        return torch.nn.Upsample(scale_factor=2, mode=upsampleMethod.lower(), align_corners=False)
 
 class Unsqueeze(torch.nn.Module):
 

konfai/network/network.py

@@ -41,57 +41,41 @@ class OptimizerLoader():
     def getOptimizer(self, key: str, parameter: Iterator[torch.nn.parameter.Parameter]) -> torch.optim.Optimizer:
         return config("{}.Model.{}.Optimizer".format(KONFAI_ROOT(), key))(getattr(importlib.import_module('torch.optim'), self.name))(parameter, config = None)
         
-class SchedulerStep():
-    
-    @config(None)
-    def __init__(self, nb_step : int = 0) -> None:
-        self.nb_step = nb_step
 
 class LRSchedulersLoader():
         
-    @config("Schedulers")
-    def __init__(self, params: dict[str, SchedulerStep] = {"default:ReduceLROnPlateau" : SchedulerStep(0)}) -> None:
-        self.params = params
-
-    def getschedulers(self, key: str, optimizer: torch.optim.Optimizer) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
-        schedulers : dict[torch.optim.lr_scheduler._LRScheduler, int] = {}
-        for 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():
+    @config()
+    def __init__(self, nb_step : int = 0) -> None:
+        self.nb_step = nb_step
+
+    def getschedulers(self, key: str, scheduler_classname: str, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
+        for m in ["torch.optim.lr_scheduler", "konfai.metric.schedulers"]:
+            module, name = _getModule(scheduler_classname, m)
+            if hasattr(importlib.import_module(module), name):
+                return config("{}.Model.{}.schedulers.{}".format(KONFAI_ROOT(), key, scheduler_classname))(getattr(importlib.import_module(module), name))(optimizer, config = None)
+        raise TrainerError("Unknown scheduler {}, tried importing from: 'torch.optim.lr_scheduler' and 'konfai.metric.schedulers', but no valid match was found. Check your YAML config or scheduler name spelling.".format(scheduler_classname))
+
+class LossSchedulersLoader():
         
-    @config("Schedulers")
-    def __init__(self, params: dict[str, SchedulerStep] = {"default:Constant" : SchedulerStep(0)}) -> None:
-        self.params = params
-
-    def getschedulers(self, key: str) -> dict[torch.optim.lr_scheduler._LRScheduler, int]:
-        schedulers : dict[Scheduler, int] = {}
-        for name, step in self.params.items():
-            if name:    
-                schedulers[getattr(importlib.import_module("konfai.metric.schedulers"), name)(config = None, DL_args = key)] = step.nb_step
-        return schedulers
+    @config()
+    def __init__(self, nb_step : int = 0) -> None:
+        self.nb_step = nb_step
+
+    def getschedulers(self, key: str, scheduler_classname: str) -> torch.optim.lr_scheduler._LRScheduler:
+        return config("{}.{}".format(key, scheduler_classname))(getattr(importlib.import_module('konfai.metric.schedulers'), scheduler_classname))(config = None)
     
 class CriterionsAttr():
     
     @config()
-    def __init__(self, l: SchedulersLoader = SchedulersLoader(), isLoss: bool = True, group: int = 0, stepStart: int = 0, stepStop: Union[int, None] = None, accumulation: bool = False) -> None:
-        self.l = l
+    def __init__(self, schedulers: dict[str, LossSchedulersLoader] = {"default:Constant": LossSchedulersLoader(0)}, isLoss: bool = True, group: int = 0, stepStart: int = 0, stepStop: Union[int, None] = None, accumulation: bool = False) -> None:
+        self.schedulersLoader = schedulers
         self.isTorchCriterion = True
         self.isLoss = isLoss
         self.stepStart = stepStart
         self.stepStop = stepStop
         self.group = group
         self.accumulation = accumulation
-        self.sheduler = None
+        self.schedulers: dict[Scheduler, int] = {}
         
 class CriterionsLoader():
 
@@ -104,7 +88,8 @@ class CriterionsLoader():
         for module_classpath, criterionsAttr in self.criterionsLoader.items():
             module, name = _getModule(module_classpath, "konfai.metric.measure")
             criterionsAttr.isTorchCriterion = module.startswith("torch")
-            criterionsAttr.sheduler = criterionsAttr.l.getschedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))
+            for scheduler_classname, schedulers in criterionsAttr.schedulersLoader.items():
+                criterionsAttr.schedulers[schedulers.getschedulers("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}.schedulers".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath), scheduler_classname)] = schedulers.nb_step
             criterions[config("{}.Model.{}.outputsCriterions.{}.targetsCriterions.{}.criterionsLoader.{}".format(KONFAI_ROOT(), model_classname, output_group, target_group, module_classpath))(getattr(importlib.import_module(module), name))(config = None)] = criterionsAttr
         return criterions
 
@@ -202,7 +187,7 @@ class Measure():
 
             for criterion, criterionsAttr in self.outputsCriterions[output_group][target_group].items():
                 if it >= criterionsAttr.stepStart and (criterionsAttr.stepStop is None or it <= criterionsAttr.stepStop):
-                    scheduler = self.update_scheduler(criterionsAttr.sheduler, it)
+                    scheduler = self.update_scheduler(criterionsAttr.schedulers, it)
                     self._loss[criterionsAttr.group]["{}:{}:{}".format(output_group, target_group, criterion.__class__.__name__)].add(scheduler.get_value(), criterion(output, *target))
                     if training and len(np.unique([len(l) for l in self._loss[criterionsAttr.group].values() if l.accumulation and l.isLoss])) == 1:
                         if criterionsAttr.isLoss:
@@ -542,7 +527,7 @@ class Network(ModuleArgsDict, ABC):
     def __init__(   self,
                     in_channels : int = 1,
                     optimizer: Union[OptimizerLoader, None] = None, 
-                    schedulers: Union[LRSchedulersLoader, None] = None, 
+                    schedulers: Union[dict[str, LRSchedulersLoader], None] = None, 
                     outputsCriterions: Union[dict[str, TargetCriterionsLoader], None] = None,
                     patch : Union[ModelPatch, None] = None,
                     nb_batch_per_step : int = 1,
@@ -556,7 +541,7 @@ class Network(ModuleArgsDict, ABC):
         self.optimizer : Union[torch.optim.Optimizer, None] = None
 
         self.LRSchedulersLoader  = schedulers
-        self.schedulers : Union[dict[torch.optim.lr_scheduler._LRScheduler, int], None] = None
+        self.schedulers : Union[dict[torch.optim.lr_scheduler._LRScheduler, int], None] = {}
 
         self.outputsCriterionsLoader = outputsCriterions
         self.measure : Union[Measure, None] = None
@@ -669,8 +654,8 @@ class Network(ModuleArgsDict, ABC):
             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:
+        for scheduler in self.schedulers:
+            if scheduler.last_epoch == -1:
                 scheduler.last_epoch = self._nb_lr_update
         self.initialized()
 
@@ -744,7 +729,8 @@ class Network(ModuleArgsDict, ABC):
                 self.optimizer.zero_grad()
 
             if self.LRSchedulersLoader and self.optimizer:
-                self.schedulers = self.LRSchedulersLoader.getschedulers(key, self.optimizer)
+                for schedulers_classname, schedulers in self.LRSchedulersLoader.items():
+                    self.schedulers[schedulers.getschedulers(key, schedulers_classname, self.optimizer)] = schedulers.nb_step
     
     def initialized(self):
         pass
@@ -892,6 +878,7 @@ class Network(ModuleArgsDict, ABC):
                 model._requires_grad(list(self.measure.outputsCriterions.keys()))
                 for loss in self.measure.getLoss():
                     self.scaler.scale(loss / self.nb_batch_per_step).backward()
+                    
                     if self._it % self.nb_batch_per_step == 0:
                         self.scaler.step(self.optimizer)
                         self.scaler.update()
@@ -903,11 +890,10 @@ class Network(ModuleArgsDict, ABC):
         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._nb_lr_update >= step and self._nb_lr_update < step+value):
-                    break
-                step += value
+        for scheduler, value in self.schedulers.items():
+            if value is None or (self._nb_lr_update >= step and self._nb_lr_update < step+value):
+                break
+            step += value
         if scheduler:
             if scheduler.__class__.__name__ == 'ReduceLROnPlateau':
                 if self.measure: