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

konfai/models/classification/resnet.py

@@ -1,115 +1,329 @@
 from abc import ABC
-from typing import Type
+
 import torch
-from konfai.network import network, blocks
-from konfai.utils.config import config
+
 from konfai.data.patching import ModelPatch
+from konfai.network import blocks, network
 
 """
-'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', dim = 2, in_channels = 3, depths=[2, 2, 2, 2], widths = [64, 64, 128, 256, 512],  num_classes=1000, useBottleneck=False
+'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+dim = 2, in_channels = 3, depths=[2, 2, 2, 2], widths = [64, 64, 128, 256, 512],
+num_classes=1000, use_bottleneck=False
 
-'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 64, 128, 256, 512],  num_classes=1000, useBottleneck=False
+'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 64, 128, 256, 512],
+num_classes=1000, use_bottleneck=False
 
-'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 256, 512, 1024, 2048],
+num_classes=1000, use_bottleneck=True
 
-'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', dim = 2, in_channels = 3, depths=[3, 4, 23, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+dim = 2, in_channels = 3, depths=[3, 4, 23, 3], widths = [64, 256, 512, 1024, 2048],
+num_classes=1000, use_bottleneck=True
 
-'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', dim = 2, in_channels = 3, depths=[3, 8, 36, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+dim = 2, in_channels = 3, depths=[3, 8, 36, 3], widths = [64, 256, 512, 1024, 2048],
+num_classes=1000, use_bottleneck=True
 
 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
 'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
 'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
 """
+
+
 class AbstractResBlock(network.ModuleArgsDict, ABC):
 
     def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
         super().__init__()
 
+
 class ResBlock(AbstractResBlock):
 
     def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
         super().__init__(in_channels, out_channels, downsample, dim)
         if downsample:
-            self.add_module("Shortcut", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=2, padding=0, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["0"], ["1"], []]), in_branch=[1], out_branch=[1], alias=["downsample"])     
-                   
-        self.add_module("ConvBlock_0", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=3, stride=2 if downsample else 1, padding=1, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
-        self.add_module("ConvBlock_1", blocks.ConvBlock(out_channels, out_channels, 1, blocks.BlockConfig(kernel_size=3, stride=1, padding=1, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv2"], ["bn2"], []]))
-        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+            self.add_module(
+                "Shortcut",
+                blocks.ConvBlock(
+                    in_channels,
+                    out_channels,
+                    blocks.BlockConfig(
+                        kernel_size=1,
+                        stride=2,
+                        padding=0,
+                        bias=False,
+                        activation="None",
+                        norm_mode=blocks.NormMode.BATCH.name,
+                    ),
+                    dim=dim,
+                    alias=[["0"], ["1"], []],
+                ),
+                in_branch=[1],
+                out_branch=[1],
+                alias=["downsample"],
+            )
+
+        self.add_module(
+            "ConvBlock_0",
+            blocks.ConvBlock(
+                in_channels,
+                out_channels,
+                blocks.BlockConfig(
+                    kernel_size=3,
+                    stride=2 if downsample else 1,
+                    padding=1,
+                    bias=False,
+                    activation="ReLU",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv1"], ["bn1"], []],
+            ),
+        )
+        self.add_module(
+            "ConvBlock_1",
+            blocks.ConvBlock(
+                out_channels,
+                out_channels,
+                blocks.BlockConfig(
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    bias=False,
+                    activation="None",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv2"], ["bn2"], []],
+            ),
+        )
+        self.add_module("Residual", blocks.Add(), in_branch=[0, 1])
         self.add_module("ReLU", torch.nn.ReLU())
 
+
 class ResBottleneckBlock(AbstractResBlock):
     def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
         super().__init__(in_channels, out_channels, downsample, dim)
-        self.add_module("ConvBlock_0", blocks.ConvBlock(in_channels, out_channels//4, 1, blocks.BlockConfig(kernel_size=1, stride=1, padding=0, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
-        self.add_module("ConvBlock_1", blocks.ConvBlock(out_channels//4, out_channels//4, 1, blocks.BlockConfig(kernel_size=3, stride=2 if downsample else 1, padding=1, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv2"], ["bn2"], []]))
-        self.add_module("ConvBlock_2", blocks.ConvBlock(out_channels//4, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=1, padding=0, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv3"], ["bn3"], []]))
-        
+        self.add_module(
+            "ConvBlock_0",
+            blocks.ConvBlock(
+                in_channels,
+                out_channels // 4,
+                blocks.BlockConfig(
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False,
+                    activation="ReLU",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv1"], ["bn1"], []],
+            ),
+        )
+        self.add_module(
+            "ConvBlock_1",
+            blocks.ConvBlock(
+                out_channels // 4,
+                out_channels // 4,
+                blocks.BlockConfig(
+                    kernel_size=3,
+                    stride=2 if downsample else 1,
+                    padding=1,
+                    bias=False,
+                    activation="ReLU",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv2"], ["bn2"], []],
+            ),
+        )
+        self.add_module(
+            "ConvBlock_2",
+            blocks.ConvBlock(
+                out_channels // 4,
+                out_channels,
+                blocks.BlockConfig(
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False,
+                    activation="ReLU",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv3"], ["bn3"], []],
+            ),
+        )
+
         if downsample or in_channels != out_channels:
-            self.add_module("Shortcut", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=2 if downsample else 1, padding=0, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["0"], ["1"], []]), in_branch=[1], out_branch=[1], alias=["downsample"])
-        
-        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+            self.add_module(
+                "Shortcut",
+                blocks.ConvBlock(
+                    in_channels,
+                    out_channels,
+                    blocks.BlockConfig(
+                        kernel_size=1,
+                        stride=2 if downsample else 1,
+                        padding=0,
+                        bias=False,
+                        activation="None",
+                        norm_mode=blocks.NormMode.BATCH.name,
+                    ),
+                    dim=dim,
+                    alias=[["0"], ["1"], []],
+                ),
+                in_branch=[1],
+                out_branch=[1],
+                alias=["downsample"],
+            )
+
+        self.add_module("Residual", blocks.Add(), in_branch=[0, 1])
         self.add_module("ReLU", torch.nn.ReLU())
 
+
 class ResNetStage(network.ModuleArgsDict):
 
-    def __init__(self, 
-                    in_channels: int,
-                    out_channels: int,
-                    depth: int,
-                    block : Type[AbstractResBlock],
-                    downsample : bool,
-                    dim : int):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        depth: int,
+        block: type[AbstractResBlock],
+        downsample: bool,
+        dim: int,
+    ):
         super().__init__()
-        self.add_module("BottleNeckBlock_0", block(in_channels=in_channels, out_channels=out_channels, downsample=downsample, dim=dim), alias=["0"])
+        self.add_module(
+            "BottleNeckBlock_0",
+            block(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                downsample=downsample,
+                dim=dim,
+            ),
+            alias=["0"],
+        )
         for i in range(1, depth):
-            self.add_module("BottleNeckBlock_{}".format(i), block(in_channels=out_channels, out_channels=out_channels, downsample=False, dim=dim), alias=["{}".format(i)])
+            self.add_module(
+                f"BottleNeckBlock_{i}",
+                block(
+                    in_channels=out_channels,
+                    out_channels=out_channels,
+                    downsample=False,
+                    dim=dim,
+                ),
+                alias=[f"{i}"],
+            )
+
 
 class ResNetStem(network.ModuleArgsDict):
 
     def __init__(self, in_channels: int, out_features: int, dim: int):
         super().__init__()
-        self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_features, 1, blocks.BlockConfig(kernel_size=7, stride=2, padding=3, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
-        self.add_module("MaxPool", blocks.getTorchModule("MaxPool", dim)(kernel_size=3, stride=2, padding=1)) 
+        self.add_module(
+            "ConvBlock",
+            blocks.ConvBlock(
+                in_channels,
+                out_features,
+                blocks.BlockConfig(
+                    kernel_size=7,
+                    stride=2,
+                    padding=3,
+                    bias=False,
+                    activation="ReLU",
+                    norm_mode=blocks.NormMode.BATCH.name,
+                ),
+                dim=dim,
+                alias=[["conv1"], ["bn1"], []],
+            ),
+        )
+        self.add_module(
+            "MaxPool",
+            blocks.get_torch_module("MaxPool", dim)(kernel_size=3, stride=2, padding=1),
+        )
+
 
 class ResNetEncoder(network.ModuleArgsDict):
-    
-    def __init__(self,
-                    in_channels: int,
-                    depths: list[int],
-                    widths: list[int],
-                    useBottleneck : bool,
-                    dim : int):
+
+    def __init__(
+        self,
+        in_channels: int,
+        depths: list[int],
+        widths: list[int],
+        use_bottleneck: bool,
+        dim: int,
+    ):
         super().__init__()
         self.add_module("ResNetStem", ResNetStem(in_channels, widths[0], dim=dim))
 
         for i, (in_channels, out_channels, depth) in enumerate(list(zip(widths[:], widths[1:], depths))):
-            self.add_module("ResNetStage_{}".format(i), ResNetStage(in_channels=in_channels, out_channels=out_channels, depth=depth, block=ResBottleneckBlock if useBottleneck else ResBlock, downsample=i!=0, dim=dim), alias=["layer{}".format(i+1)])
+            self.add_module(
+                f"ResNetStage_{i}",
+                ResNetStage(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    depth=depth,
+                    block=ResBottleneckBlock if use_bottleneck else ResBlock,
+                    downsample=i != 0,
+                    dim=dim,
+                ),
+                alias=[f"layer{i + 1}"],
+            )
+
 
 class Head(network.ModuleArgsDict):
-    
-    def __init__(self, in_features : int, num_classes : int, dim : int):
+
+    def __init__(self, in_features: int, num_classes: int, dim: int):
         super().__init__()
-        self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(1))
+        self.add_module("AdaptiveAvgPool", blocks.get_torch_module("AdaptiveAvgPool", dim)(1))
         self.add_module("Flatten", torch.nn.Flatten(1))
-        self.add_module("Linear", torch.nn.Linear(in_features, num_classes), pretrained=False, alias=["fc"])
+        self.add_module(
+            "Linear",
+            torch.nn.Linear(in_features, num_classes),
+            pretrained=False,
+            alias=["fc"],
+        )
         self.add_module("Unsqueeze", blocks.Unsqueeze(2))
 
-        
+
 class ResNet(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] = [2, 2, 2, 2],
-                    widths: list[int] = [64, 64, 128, 256, 512],
-                    num_classes: int = 10, 
-                    useBottleneck=False):
-        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("ResNetEncoder", ResNetEncoder(in_channels=in_channels, depths=depths, widths=widths, useBottleneck=useBottleneck, 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] = [2, 2, 2, 2],
+        widths: list[int] = [64, 64, 128, 256, 512],
+        num_classes: int = 10,
+        use_bottleneck=False,
+    ):
+        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(
+            "ResNetEncoder",
+            ResNetEncoder(
+                in_channels=in_channels,
+                depths=depths,
+                widths=widths,
+                use_bottleneck=use_bottleneck,
+                dim=dim,
+            ),
+        )
         self.add_module("Head", Head(in_features=widths[-1], num_classes=num_classes, dim=dim))
-