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

konfai/metric/schedulers.py

@@ -1,15 +1,16 @@
-import torch
-import numpy as np
 from abc import abstractmethod
-from konfai.utils.config import config
 from functools import partial
 
-class Scheduler():
+import numpy as np
+import torch
+
+
+class Scheduler:
 
     def __init__(self, start_value: float) -> None:
         self.baseValue = float(start_value)
         self.it = 0
-        
+
     def step(self, it: int):
         self.it = it
 
@@ -17,6 +18,7 @@ class Scheduler():
     def get_value(self) -> float:
         pass
 
+
 class Constant(Scheduler):
 
     def __init__(self, value: float = 1):
@@ -25,20 +27,29 @@ class Constant(Scheduler):
     def get_value(self) -> float:
         return self.baseValue
 
+
 class CosineAnnealing(Scheduler):
-    
-    def __init__(self, start_value: float = 1, eta_min: float = 0.00001, T_max: int = 100):
+
+    def __init__(self, start_value: float = 1, eta_min: float = 0.00001, t_max: int = 100):
         super().__init__(start_value)
         self.eta_min = eta_min
-        self.T_max = T_max
+        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+1) / (warmup_steps+1))
 
-    def __init__(self, optimizer: torch.optim.Optimizer, warmup_steps: int = 10, last_epoch=-1, verbose="deprecated"):
+    @staticmethod
+    def warmup(warmup_steps: int, step: int) -> float:
+        return min(1.0, (step + 1) / (warmup_steps + 1))
+
+    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/classification/convNeXt.py

@@ -1,33 +1,44 @@
 import torch
-import torch.nn.functional as F
-from konfai.network import network, blocks
-from konfai.utils.config import config
+import torch.nn.functional as F  # noqa: N812
+
 from konfai.data.patching import ModelPatch
+from konfai.network import blocks, network
 
 """
-"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth", depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]
-"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth", depths=[3, 3, 27, 3], dims=[96, 192, 384, 768]
-"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth", depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]
-"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",[3, 3, 27, 3], dims=[192, 384, 768, 1536]
+"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]
+"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+depths=[3, 3, 27, 3], dims=[96, 192, 384, 768]
+"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]
+"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+[3, 3, 27, 3], dims=[192, 384, 768, 1536]
 "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
 "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
 "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
 "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
-"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth", depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048]
+"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
+depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048]
 """
 
+
 class LayerNorm(torch.nn.Module):
 
-    def __init__(self, normalized_shape : int, eps : float = 1e-6, data_format : str="channels_last"):
+    def __init__(
+        self,
+        normalized_shape: int,
+        eps: float = 1e-6,
+        data_format: str = "channels_last",
+    ):
         super().__init__()
         self.weight = torch.nn.parameter.Parameter(torch.ones(normalized_shape))
         self.bias = torch.nn.parameter.Parameter(torch.zeros(normalized_shape))
         self.eps = eps
         self.data_format = data_format
         if self.data_format not in ["channels_last", "channels_first"]:
-            raise NotImplementedError 
-        self.normalized_shape = (normalized_shape, )
-    
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
     def forward(self, x):
         if self.data_format == "channels_last":
             return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
@@ -44,17 +55,18 @@ class LayerNorm(torch.nn.Module):
             return x
 
     def extra_repr(self):
-        return "normalized_shape={}, eps={}, data_format={})".format(self.normalized_shape, self.eps, self.data_format)
+        return f"normalized_shape={self.normalized_shape}, eps={self.eps}, data_format={self.data_format})"
+
 
 class DropPath(torch.nn.Module):
 
-    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
-        super(DropPath, self).__init__()
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super().__init__()
         self.drop_prob = drop_prob
         self.scale_by_keep = scale_by_keep
 
     def forward(self, x):
-        if self.drop_prob == 0. or not self.training:
+        if self.drop_prob == 0.0 or not self.training:
             return x
         keep_prob = 1 - self.drop_prob
         shape = (x.shape[0],) + (1,) * (x.ndim - 1)
@@ -64,111 +76,205 @@ class DropPath(torch.nn.Module):
         return x * random_tensor
 
     def extra_repr(self):
-        return "drop_prob={}".format(round(self.drop_prob,3))
+        return f"drop_prob={round(self.drop_prob, 3)}"
+
 
 class LayerScaler(torch.nn.Module):
-    
-    def __init__(self, init_value : float, dimensions : int):
+
+    def __init__(self, init_value: float, dimensions: int):
         super().__init__()
         self.init_value = init_value
         self.gamma = torch.nn.Parameter(torch.ones(dimensions, 1, 1) * init_value)
 
-    def forward(self, input : torch.Tensor) -> torch.Tensor:
-        return self.gamma * input
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self.gamma * tensor
 
     def extra_repr(self):
-        return "init_value={}".format(self.init_value)
+        return f"init_value={self.init_value}"
+
 
 class BottleNeckBlock(network.ModuleArgsDict):
-    def __init__(   self,
-                    features: int,
-                    drop_p: float,
-                    layer_scaler_init_value: float,
-                    dim: int):
+    def __init__(self, features: int, drop_p: float, layer_scaler_init_value: float, dim: int):
         super().__init__()
-        self.add_module("Conv_0", blocks.getTorchModule("Conv", dim)(features, features, kernel_size=7, padding=3, groups=features), alias=["dwconv"])
+        self.add_module(
+            "Conv_0",
+            blocks.get_torch_module("Conv", dim)(features, features, kernel_size=7, padding=3, groups=features),
+            alias=["dwconv"],
+        )
         self.add_module("ToFeatures", blocks.ToFeatures(dim))
         self.add_module("LayerNorm", LayerNorm(features, eps=1e-6), alias=["norm"])
-        self.add_module("Linear_1",  torch.nn.Linear(features, features * 4), alias=["pwconv1"])
+        self.add_module("Linear_1", torch.nn.Linear(features, features * 4), alias=["pwconv1"])
         self.add_module("GELU", torch.nn.GELU())
         self.add_module("Linear_2", torch.nn.Linear(features * 4, features), alias=["pwconv2"])
         self.add_module("ToChannels", blocks.ToChannels(dim))
-        self.add_module("LayerScaler", LayerScaler(init_value=layer_scaler_init_value, dimensions=features), alias=[""])
+        self.add_module(
+            "LayerScaler",
+            LayerScaler(init_value=layer_scaler_init_value, dimensions=features),
+            alias=[""],
+        )
         self.add_module("StochasticDepth", DropPath(drop_p))
-        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
-        
+        self.add_module("Residual", blocks.Add(), in_branch=[0, 1])
+
+
 class DownSample(network.ModuleArgsDict):
-    
-    def __init__(   self, 
-                    in_features: int,
-                    out_features: int,
-                    dim : int):
+
+    def __init__(self, in_features: int, out_features: int, dim: int):
         super().__init__()
-        self.add_module("LayerNorm", LayerNorm(in_features, eps=1e-6, data_format="channels_first"), alias=["0"])
-        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_features, out_features, kernel_size=2, stride=2), alias=["1"])
+        self.add_module(
+            "LayerNorm",
+            LayerNorm(in_features, eps=1e-6, data_format="channels_first"),
+            alias=["0"],
+        )
+        self.add_module(
+            "Conv",
+            blocks.get_torch_module("Conv", dim)(in_features, out_features, kernel_size=2, stride=2),
+            alias=["1"],
+        )
+
 
 class ConvNexStage(network.ModuleArgsDict):
-    
-    def __init__(   self, 
-                    features: int,
-                    depth: int, 
-                    drop_p: list[float],
-                    dim : int):
+
+    def __init__(self, features: int, depth: int, drop_p: list[float], dim: int):
         super().__init__()
         for i in range(depth):
-            self.add_module("BottleNeckBlock_{}".format(i), BottleNeckBlock(features=features, drop_p=drop_p[i], layer_scaler_init_value=1e-6, dim=dim), alias=["{}".format(i)])
+            self.add_module(
+                f"BottleNeckBlock_{i}",
+                BottleNeckBlock(
+                    features=features,
+                    drop_p=drop_p[i],
+                    layer_scaler_init_value=1e-6,
+                    dim=dim,
+                ),
+                alias=[f"{i}"],
+            )
+
 
 class ConvNextStem(network.ModuleArgsDict):
 
     def __init__(self, in_features: int, out_features: int, dim: int):
         super().__init__()
-        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_features, out_features, kernel_size=4, stride=4), alias=["0"])
-        self.add_module("LayerNorm", LayerNorm(out_features, eps=1e-6, data_format="channels_first"), alias=["1"])
+        self.add_module(
+            "Conv",
+            blocks.get_torch_module("Conv", dim)(in_features, out_features, kernel_size=4, stride=4),
+            alias=["0"],
+        )
+        self.add_module(
+            "LayerNorm",
+            LayerNorm(out_features, eps=1e-6, data_format="channels_first"),
+            alias=["1"],
+        )
+
 
 class ConvNextEncoder(network.ModuleArgsDict):
 
-    def __init__(   self,
-                    in_channels: int,
-                    depths: list[int],
-                    widths: list[int],
-                    drop_p: float,
-                    dim : int):
+    def __init__(
+        self,
+        in_channels: int,
+        depths: list[int],
+        widths: list[int],
+        drop_p: float,
+        dim: int,
+    ):
         super().__init__()
-        self.add_module("ConvNextStem", ConvNextStem(in_channels, widths[0], dim=dim), alias=["downsample_layers.0"])
-        
+        self.add_module(
+            "ConvNextStem",
+            ConvNextStem(in_channels, widths[0], dim=dim),
+            alias=["downsample_layers.0"],
+        )
+
         drop_probs = [x.item() for x in torch.linspace(0, drop_p, sum(depths))]
-        self.add_module("ConvNexStage_0", ConvNexStage(features=widths[0], depth=depths[0], drop_p=drop_probs[:depths[0]], dim=dim), alias=["stages.0"])
-        
+        self.add_module(
+            "ConvNexStage_0",
+            ConvNexStage(
+                features=widths[0],
+                depth=depths[0],
+                drop_p=drop_probs[: depths[0]],
+                dim=dim,
+            ),
+            alias=["stages.0"],
+        )
+
         for i, (in_features, out_features) in enumerate(list(zip(widths[:], widths[1:]))):
-            self.add_module("DownSample_{}".format(i+1), DownSample(in_features=in_features, out_features=out_features, dim=dim), alias=["downsample_layers.{}".format(i+1)])
-            self.add_module("ConvNexStage_{}".format(i+1), ConvNexStage(features=out_features, depth=depths[i+1], drop_p=drop_probs[sum(depths[:i+1]):sum(depths[:i+2])], dim=dim), alias=["stages.{}".format(i+1)])
+            self.add_module(
+                f"DownSample_{i + 1}",
+                DownSample(in_features=in_features, out_features=out_features, dim=dim),
+                alias=[f"downsample_layers.{i + 1}"],
+            )
+            self.add_module(
+                f"ConvNexStage_{i + 1}",
+                ConvNexStage(
+                    features=out_features,
+                    depth=depths[i + 1],
+                    drop_p=drop_probs[sum(depths[: i + 1]) : sum(depths[: i + 2])],
+                    dim=dim,
+                ),
+                alias=[f"stages.{i + 1}"],
+            )
+
 
 class Head(network.ModuleArgsDict):
 
-    def __init__(self, in_features : int, num_classes : list[int], dim : int) -> None:
+    def __init__(self, in_features: int, num_classes: list[int], dim: int) -> None:
         super().__init__()
-        self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(tuple([1]*dim)))
+        self.add_module(
+            "AdaptiveAvgPool",
+            blocks.get_torch_module("AdaptiveAvgPool", dim)(tuple([1] * dim)),
+        )
         self.add_module("Flatten", torch.nn.Flatten(1))
         self.add_module("LayerNorm", torch.nn.LayerNorm(in_features, eps=1e-6), alias=["norm"])
-        
+
         for i, nb_classe in enumerate(num_classes):
-            self.add_module("Linear_{}".format(i), torch.nn.Linear(in_features, nb_classe), pretrained=False, alias=["head"], out_branch=[i+1])
-            self.add_module("Unsqueeze_{}".format(i), blocks.Unsqueeze(2), in_branch=[i+1], out_branch=[-1])
+            self.add_module(
+                f"Linear_{i}",
+                torch.nn.Linear(in_features, nb_classe),
+                pretrained=False,
+                alias=["head"],
+                out_branch=[i + 1],
+            )
+            self.add_module(
+                f"Unsqueeze_{i}",
+                blocks.Unsqueeze(2),
+                in_branch=[i + 1],
+                out_branch=[-1],
+            )
+
 
 class ConvNeXt(network.Network):
-    
-    def __init__(   self,
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    patch : ModelPatch = ModelPatch(),
-                    dim : int = 3,
-                    in_channels: int = 1,
-                    depths: list[int] = [3,3,27,3],
-                    widths: list[int] = [128, 256, 512, 1024],
-                    drop_p: float = 0.1,
-                    num_classes: list[int] = [4, 7]):
-
-        super().__init__(in_channels = in_channels, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim = dim, patch=patch, init_type = "trunc_normal", init_gain=0.02)
-        self.add_module("ConvNextEncoder", ConvNextEncoder(in_channels=in_channels, depths=depths, widths=widths, drop_p=drop_p, dim=dim))        
-        self.add_module("Head", Head(in_features=widths[-1], num_classes=num_classes, dim=dim))
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        patch: ModelPatch = ModelPatch(),
+        dim: int = 3,
+        in_channels: int = 1,
+        depths: list[int] = [3, 3, 27, 3],
+        widths: list[int] = [128, 256, 512, 1024],
+        drop_p: float = 0.1,
+        num_classes: list[int] = [4, 7],
+    ):
+
+        super().__init__(
+            in_channels=in_channels,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            patch=patch,
+            init_type="trunc_normal",
+            init_gain=0.02,
+        )
+        self.add_module(
+            "ConvNextEncoder",
+            ConvNextEncoder(
+                in_channels=in_channels,
+                depths=depths,
+                widths=widths,
+                drop_p=drop_p,
+                dim=dim,
+            ),
+        )
+        self.add_module("Head", Head(in_features=widths[-1], num_classes=num_classes, dim=dim))