updates

app.py

@@ -73,843 +73,846 @@ effnet_hparams = {61: {
 activation_indices = {'efficientnet': [0, 3]}
 
 
-########## EfficientNet ############
-@dataclass
-class _MBConvConfig:
-    expand_ratio: float
-    kernel: int
-    stride: int
-    input_channels: int
-    out_channels: int
-    num_layers: int
-    block: Callable[..., nn.Module]
-
-    @staticmethod
-    def adjust_channels(
-        channels: int, width_mult: float, min_value: Optional[int] = None
-    ) -> int:
-        return _make_divisible(channels * width_mult, 8, min_value)
 
+#### extra torchvision stuff ####
 
-class MBConvConfig(_MBConvConfig):
-    # Stores information listed at Table 1 of the EfficientNet paper & Table 4 of the EfficientNetV2 paper
-    def __init__(
-        self,
-        expand_ratio: float,
-        kernel: int,
-        stride: int,
-        input_channels: int,
-        out_channels: int,
-        num_layers: int,
-        width_mult: float = 1.0,
-        depth_mult: float = 1.0,
-        block: Optional[Callable[..., nn.Module]] = None,
-    ) -> None:
-        input_channels = self.adjust_channels(input_channels, width_mult)
-        out_channels = self.adjust_channels(out_channels, width_mult)
-        num_layers = self.adjust_depth(num_layers, depth_mult)
-        if block is None:
-            block = MBConv
-        super().__init__(
-            expand_ratio,
-            kernel,
-            stride,
-            input_channels,
-            out_channels,
-            num_layers,
-            block,
-        )
 
-    @staticmethod
-    def adjust_depth(num_layers: int, depth_mult: float):
-        return int(math.ceil(num_layers * depth_mult))
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed
 
+    Args:
+        num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)``
+        eps (float): a value added to the denominator for numerical stability. Default: 1e-5
+    """
 
-class FusedMBConvConfig(_MBConvConfig):
-    # Stores information listed at Table 4 of the EfficientNetV2 paper
     def __init__(
         self,
-        expand_ratio: float,
-        kernel: int,
-        stride: int,
-        input_channels: int,
-        out_channels: int,
-        num_layers: int,
-        block: Optional[Callable[..., nn.Module]] = None,
-    ) -> None:
-        if block is None:
-            block = FusedMBConv
-        super().__init__(
-            expand_ratio,
-            kernel,
-            stride,
-            input_channels,
-            out_channels,
-            num_layers,
-            block,
+        num_features: int,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        # _log_api_usage_once(self)
+        self.eps = eps
+        self.register_buffer("weight", torch.ones(num_features))
+        self.register_buffer("bias", torch.zeros(num_features))
+        self.register_buffer("running_mean", torch.zeros(num_features))
+        self.register_buffer("running_var", torch.ones(num_features))
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict,
+        prefix: str,
+        local_metadata: dict,
+        strict: bool,
+        missing_keys: List[str],
+        unexpected_keys: List[str],
+        error_msgs: List[str],
+    ):
+        num_batches_tracked_key = prefix + "num_batches_tracked"
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
         )
 
+    def forward(self, x: Tensor) -> Tensor:
+        # move reshapes to the beginning
+        # to make it fuser-friendly
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        scale = w * (rv + self.eps).rsqrt()
+        bias = b - rm * scale
+        return x * scale + bias
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})"
 
-class MBConv(nn.Module):
+
+class ConvNormActivation(torch.nn.Sequential):
     def __init__(
         self,
-        cnf: MBConvConfig,
-        stochastic_depth_prob: float,
-        norm_layer: Callable[..., nn.Module],
-        se_layer: Callable[..., nn.Module] = SqueezeExcitation,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: Optional[int] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: int = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+        conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d,
     ) -> None:
-        super().__init__()
-
-        if not (1 <= cnf.stride <= 2):
-            raise ValueError("illegal stride value")
-
-        self.use_res_connect = (
-            cnf.stride == 1 and cnf.input_channels == cnf.out_channels
-        )
 
-        layers: List[nn.Module] = []
-        activation_layer = nn.SiLU
+        if padding is None:
+            padding = (kernel_size - 1) // 2 * dilation
+        if bias is None:
+            bias = norm_layer is None
 
-        # expand
-        expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio)
-        if expanded_channels != cnf.input_channels:
-            layers.append(
-                Conv2dNormActivation(
-                    cnf.input_channels,
-                    expanded_channels,
-                    kernel_size=1,
-                    norm_layer=norm_layer,
-                    activation_layer=activation_layer,
-                )
+        layers = [
+            conv_layer(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
             )
+        ]
 
-        # depthwise
-        layers.append(
-            Conv2dNormActivation(
-                expanded_channels,
-                expanded_channels,
-                kernel_size=cnf.kernel,
-                stride=cnf.stride,
-                groups=expanded_channels,
-                norm_layer=norm_layer,
-                activation_layer=activation_layer,
-            )
-        )
+        if norm_layer is not None:
+            layers.append(norm_layer(out_channels))
 
-        # squeeze and excitation
-        squeeze_channels = max(1, cnf.input_channels // 4)
-        layers.append(
-            se_layer(
-                expanded_channels,
-                squeeze_channels,
-                activation=partial(nn.SiLU, inplace=True),
-            )
-        )
+        if activation_layer is not None:
+            params = {} if inplace is None else {"inplace": inplace}
+            layers.append(activation_layer(**params))
+        super().__init__(*layers)
+        # _log_api_usage_once(self)
+        self.out_channels = out_channels
 
-        # project
-        layers.append(
-            Conv2dNormActivation(
-                expanded_channels,
-                cnf.out_channels,
-                kernel_size=1,
-                norm_layer=norm_layer,
-                activation_layer=None,
+        if self.__class__ == ConvNormActivation:
+            warnings.warn(
+                "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead."
             )
-        )
 
-        self.block = nn.Sequential(*layers)
-        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
-        self.out_channels = cnf.out_channels
 
-    def forward(self, input: Tensor) -> Tensor:
-        result = self.block(input)
-        if self.use_res_connect:
-            result = self.stochastic_depth(result)
-            result += input
-        return result
+class Conv2dNormActivation(ConvNormActivation):
+    """
+    Configurable block used for Convolution2d-Normalization-Activation blocks.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
+        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm2d``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
+        dilation (int): Spacing between kernel elements. Default: 1
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
 
+    """
 
-class FusedMBConv(nn.Module):
     def __init__(
         self,
-        cnf: FusedMBConvConfig,
-        stochastic_depth_prob: float,
-        norm_layer: Callable[..., nn.Module],
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: Optional[int] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: int = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
     ) -> None:
-        super().__init__()
-
-        if not (1 <= cnf.stride <= 2):
-            raise ValueError("illegal stride value")
 
-        self.use_res_connect = (
-            cnf.stride == 1 and cnf.input_channels == cnf.out_channels
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            groups,
+            norm_layer,
+            activation_layer,
+            dilation,
+            inplace,
+            bias,
+            torch.nn.Conv2d,
         )
 
-        layers: List[nn.Module] = []
-        activation_layer = nn.SiLU
-
-        expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio)
-        if expanded_channels != cnf.input_channels:
-            # fused expand
-            layers.append(
-                Conv2dNormActivation(
-                    cnf.input_channels,
-                    expanded_channels,
-                    kernel_size=cnf.kernel,
-                    stride=cnf.stride,
-                    norm_layer=norm_layer,
-                    activation_layer=activation_layer,
-                )
-            )
-
-            # project
-            layers.append(
-                Conv2dNormActivation(
-                    expanded_channels,
-                    cnf.out_channels,
-                    kernel_size=1,
-                    norm_layer=norm_layer,
-                    activation_layer=None,
-                )
-            )
-        else:
-            layers.append(
-                Conv2dNormActivation(
-                    cnf.input_channels,
-                    cnf.out_channels,
-                    kernel_size=cnf.kernel,
-                    stride=cnf.stride,
-                    norm_layer=norm_layer,
-                    activation_layer=activation_layer,
-                )
-            )
-
-        self.block = nn.Sequential(*layers)
-        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
-        self.out_channels = cnf.out_channels
 
-    def forward(self, input: Tensor) -> Tensor:
-        result = self.block(input)
-        if self.use_res_connect:
-            result = self.stochastic_depth(result)
-            result += input
-        return result
+class Conv3dNormActivation(ConvNormActivation):
+    """
+    Configurable block used for Convolution3d-Normalization-Activation blocks.
 
+    Args:
+        in_channels (int): Number of channels in the input video.
+        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
+        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm3d``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
+        dilation (int): Spacing between kernel elements. Default: 1
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+    """
 
-class EfficientNetConfig(PretrainedConfig):
-    
-    model_type = "efficientnet"
-    
     def __init__(
         self,
-        # inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
-        dropout: float=0.25,
-        num_channels: int = 61,
-        stochastic_depth_prob: float = 0.2,
-        num_classes: int = 2,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
-        # last_channel: Optional[int] = None,
-        size: str='v2_s',
-        width_mult: float = 1.0,
-        depth_mult: float = 1.0,
-        **kwargs: Any,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: Optional[int] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm3d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: int = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
     ) -> None:
-        """
-        EfficientNet V1 and V2 main class
 
-        Args:
-            inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure
-            dropout (float): The droupout probability
-            stochastic_depth_prob (float): The stochastic depth probability
-            num_classes (int): Number of classes
-            norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use
-            last_channel (int): The number of channels on the penultimate layer
-        """
-        
-        
-        # self.model = EfficientNet(
-        #                             dropout=dropout,
-        #                             num_channels=num_channels,
-        #                             num_classes=num_classes,
-        #                             size=size,
-        #                             stochastic_depth_prob=stochastic_depth_prob,
-        #                             width_mult=width_mult,
-        #                             depth_mult=depth_mult,
-        # )
-        
-        # 
-        self.dropout=dropout
-        self.num_channels=num_channels
-        self.num_classes=num_classes
-        self.size=size
-        self.stochastic_depth_prob=stochastic_depth_prob
-        self.width_mult=width_mult
-        self.depth_mult=depth_mult
-                                    
-        super().__init__(**kwargs)
-        
-    
-class EfficientNetPreTrained(PreTrainedModel):
-    
-    config_class = EfficientNetConfig
-    
-    def __init__(
-        self,
-        config
-    ):
-        super().__init__(config)   
-        self.model = EfficientNet(  dropout=config.dropout,
-                                    num_channels=config.num_channels,
-                                    num_classes=config.num_classes,
-                                    size=config.size,
-                                    stochastic_depth_prob=config.stochastic_depth_prob,
-                                    width_mult=config.width_mult,
-                                    depth_mult=config.depth_mult,)
-        
-    def forward(self, tensor):
-        return self.model.forward(tensor)
-    
-    
-class EfficientNet(nn.Module):
-    
-    
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            groups,
+            norm_layer,
+            activation_layer,
+            dilation,
+            inplace,
+            bias,
+            torch.nn.Conv3d,
+        )
+
+
+class SqueezeExcitation(torch.nn.Module):
+    """
+    This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1).
+    Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in eq. 3.
+
+    Args:
+        input_channels (int): Number of channels in the input image
+        squeeze_channels (int): Number of squeeze channels
+        activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU``
+        scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid``
+    """
+
     def __init__(
         self,
-        # inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
-        dropout: float=0.25,
-        num_channels: int = 61,
-        stochastic_depth_prob: float = 0.2,
-        num_classes: int = 2,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
-        # last_channel: Optional[int] = None,
-        size: str='v2_s',
-        width_mult: float = 1.0,
-        depth_mult: float = 1.0,
-        **kwargs: Any,
+        input_channels: int,
+        squeeze_channels: int,
+        activation: Callable[..., torch.nn.Module] = torch.nn.ReLU,
+        scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid,
     ) -> None:
-        """
-        EfficientNet V1 and V2 main class
-
-        Args:
-            inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure
-            dropout (float): The droupout probability
-            stochastic_depth_prob (float): The stochastic depth probability
-            num_classes (int): Number of classes
-            norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use
-            last_channel (int): The number of channels on the penultimate layer
-        """
         super().__init__()
         # _log_api_usage_once(self)
-        
-        inverted_residual_setting, last_channel = _efficientnet_conf(
-                     "efficientnet_%s" % (size), width_mult=width_mult, depth_mult=depth_mult
-                    )
+        self.avgpool = torch.nn.AdaptiveAvgPool2d(1)
+        self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1)
+        self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1)
+        self.activation = activation()
+        self.scale_activation = scale_activation()
 
-        if not inverted_residual_setting:
-            raise ValueError("The inverted_residual_setting should not be empty")
-        elif not (
-            isinstance(inverted_residual_setting, Sequence)
-            and all([isinstance(s, _MBConvConfig) for s in inverted_residual_setting])
-        ):
-            raise TypeError(
-                "The inverted_residual_setting should be List[MBConvConfig]"
-            )
+    def _scale(self, input: Tensor) -> Tensor:
+        scale = self.avgpool(input)
+        scale = self.fc1(scale)
+        scale = self.activation(scale)
+        scale = self.fc2(scale)
+        return self.scale_activation(scale)
 
-        if "block" in kwargs:
-            warnings.warn(
-                "The parameter 'block' is deprecated since 0.13 and will be removed 0.15. "
-                "Please pass this information on 'MBConvConfig.block' instead."
-            )
-            if kwargs["block"] is not None:
-                for s in inverted_residual_setting:
-                    if isinstance(s, MBConvConfig):
-                        s.block = kwargs["block"]
+    def forward(self, input: Tensor) -> Tensor:
+        scale = self._scale(input)
+        return scale * input
 
-        if norm_layer is None:
-            norm_layer = nn.BatchNorm2d
 
-        layers: List[nn.Module] = []
-
-        # building first layer
-        firstconv_output_channels = inverted_residual_setting[0].input_channels
-        layers.append(
-            Conv2dNormActivation(
-                num_channels,
-                firstconv_output_channels,
-                kernel_size=3,
-                stride=2,
-                norm_layer=norm_layer,
-                activation_layer=nn.SiLU,
-            )
-        )
-
-        # building inverted residual blocks
-        total_stage_blocks = sum(cnf.num_layers for cnf in inverted_residual_setting)
-        stage_block_id = 0
-        for cnf in inverted_residual_setting:
-            stage: List[nn.Module] = []
-            for _ in range(cnf.num_layers):
-                # copy to avoid modifications. shallow copy is enough
-                block_cnf = copy.copy(cnf)
-
-                # overwrite info if not the first conv in the stage
-                if stage:
-                    block_cnf.input_channels = block_cnf.out_channels
-                    block_cnf.stride = 1
-
-                # adjust stochastic depth probability based on the depth of the stage block
-                sd_prob = (
-                    stochastic_depth_prob * float(stage_block_id) / total_stage_blocks
-                )
+class MLP(torch.nn.Sequential):
+    """This block implements the multi-layer perceptron (MLP) module.
 
-                stage.append(block_cnf.block(block_cnf, sd_prob, norm_layer))
-                stage_block_id += 1
+    Args:
+        in_channels (int): Number of channels of the input
+        hidden_channels (List[int]): List of the hidden channel dimensions
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``None``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool): Whether to use bias in the linear layer. Default ``True``
+        dropout (float): The probability for the dropout layer. Default: 0.0
+    """
 
-            layers.append(nn.Sequential(*stage))
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: List[int],
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = None,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        inplace: Optional[bool] = True,
+        bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal:
+        # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py
+        params = {} if inplace is None else {"inplace": inplace}
 
-        # building last several layers
-        lastconv_input_channels = inverted_residual_setting[-1].out_channels
-        lastconv_output_channels = (
-            last_channel if last_channel is not None else 4 * lastconv_input_channels
-        )
-        layers.append(
-            Conv2dNormActivation(
-                lastconv_input_channels,
-                lastconv_output_channels,
-                kernel_size=1,
-                norm_layer=norm_layer,
-                activation_layer=nn.SiLU,
-            )
-        )
+        layers = []
+        in_dim = in_channels
+        for hidden_dim in hidden_channels[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
+            if norm_layer is not None:
+                layers.append(norm_layer(hidden_dim))
+            layers.append(activation_layer(**params))
+            layers.append(torch.nn.Dropout(dropout, **params))
+            in_dim = hidden_dim
 
-        self.features = nn.Sequential(*layers)
-        self.avgpool = nn.AdaptiveAvgPool2d(1)
-        self.classifier = nn.Sequential(
-            nn.Dropout(p=dropout, inplace=True),
-            nn.Linear(lastconv_output_channels, num_classes),
-        )
+        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias))
+        layers.append(torch.nn.Dropout(dropout, **params))
 
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode="fan_out")
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
-                nn.init.ones_(m.weight)
-                nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.Linear):
-                init_range = 1.0 / math.sqrt(m.out_features)
-                nn.init.uniform_(m.weight, -init_range, init_range)
-                nn.init.zeros_(m.bias)
-                
-        # super().__init__(**kwargs)
+        super().__init__(*layers)
+        # _log_api_usage_once(self)
 
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
 
-        x = self.avgpool(x)
-        x = torch.flatten(x, 1)
+class Permute(torch.nn.Module):
+    """This module returns a view of the tensor input with its dimensions permuted.
 
-        x = self.classifier(x)
+    Args:
+        dims (List[int]): The desired ordering of dimensions
+    """
 
-        return x
+    def __init__(self, dims: List[int]):
+        super().__init__()
+        self.dims = dims
 
     def forward(self, x: Tensor) -> Tensor:
-        return self._forward_impl(x)
-
-
-# def _efficientnet(
-#     inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
-#     dropout: float,
-#     last_channel: Optional[int],
-#     weights=None,
-#     num_channels: int = 61,
-#     stochastic_depth_prob: float = 0.2,
-#     progress: bool = True,
-#     num_classes: int = 2,
-#     **kwargs: Any,
-# ) -> EfficientNetCongig:
-
-#     model = EfficientNetCongif(
-#         inverted_residual_setting,
-#         dropout,
-#         num_classes=num_classes,
-#         num_channels=num_channels,
-#         stochastic_depth_prob=stochastic_depth_prob,
-#         last_channel=last_channel,
-#         **kwargs,
-#     )
+        return torch.permute(x, self.dims)
 
-#     return model
 
 
-def _efficientnet_conf(
-    arch: str,
-    **kwargs: Any,
-) -> Tuple[Sequence[Union[MBConvConfig, FusedMBConvConfig]], Optional[int]]:
-    inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]]
-    if arch.startswith("efficientnet_b"):
-        bneck_conf = partial(
-            MBConvConfig,
-            width_mult=kwargs.pop("width_mult"),
-            depth_mult=kwargs.pop("depth_mult"),
-        )
-        inverted_residual_setting = [
-            bneck_conf(1, 3, 1, 32, 16, 1),
-            bneck_conf(6, 3, 2, 16, 24, 2),
-            bneck_conf(6, 5, 2, 24, 40, 2),
-            bneck_conf(6, 3, 2, 40, 80, 3),
-            bneck_conf(6, 5, 1, 80, 112, 3),
-            bneck_conf(6, 5, 2, 112, 192, 4),
-            bneck_conf(6, 3, 1, 192, 320, 1),
-        ]
-        last_channel = None
-    elif arch.startswith("efficientnet_v2_s"):
-        inverted_residual_setting = [
-            FusedMBConvConfig(1, 3, 1, 24, 24, 2),
-            FusedMBConvConfig(4, 3, 2, 24, 48, 4),
-            FusedMBConvConfig(4, 3, 2, 48, 64, 4),
-            MBConvConfig(4, 3, 2, 64, 128, 6),
-            MBConvConfig(6, 3, 1, 128, 160, 9),
-            MBConvConfig(6, 3, 2, 160, 256, 15),
-        ]
-        last_channel = 1280
-    elif arch.startswith("efficientnet_v2_m"):
-        inverted_residual_setting = [
-            FusedMBConvConfig(1, 3, 1, 24, 24, 3),
-            FusedMBConvConfig(4, 3, 2, 24, 48, 5),
-            FusedMBConvConfig(4, 3, 2, 48, 80, 5),
-            MBConvConfig(4, 3, 2, 80, 160, 7),
-            MBConvConfig(6, 3, 1, 160, 176, 14),
-            MBConvConfig(6, 3, 2, 176, 304, 18),
-            MBConvConfig(6, 3, 1, 304, 512, 5),
-        ]
-        last_channel = 1280
-    elif arch.startswith("efficientnet_v2_l"):
-        inverted_residual_setting = [
-            FusedMBConvConfig(1, 3, 1, 32, 32, 4),
-            FusedMBConvConfig(4, 3, 2, 32, 64, 7),
-            FusedMBConvConfig(4, 3, 2, 64, 96, 7),
-            MBConvConfig(4, 3, 2, 96, 192, 10),
-            MBConvConfig(6, 3, 1, 192, 224, 19),
-            MBConvConfig(6, 3, 2, 224, 384, 25),
-            MBConvConfig(6, 3, 1, 384, 640, 7),
-        ]
-        last_channel = 1280
-    else:
-        raise ValueError(f"Unsupported model type {arch}")
 
-    return inverted_residual_setting, last_channel
+########## EfficientNet ############
+@dataclass
+class _MBConvConfig:
+    expand_ratio: float
+    kernel: int
+    stride: int
+    input_channels: int
+    out_channels: int
+    num_layers: int
+    block: Callable[..., nn.Module]
 
+    @staticmethod
+    def adjust_channels(
+        channels: int, width_mult: float, min_value: Optional[int] = None
+    ) -> int:
+        return _make_divisible(channels * width_mult, 8, min_value)
 
-#### extra torchvision stuff ####
 
+class MBConvConfig(_MBConvConfig):
+    # Stores information listed at Table 1 of the EfficientNet paper & Table 4 of the EfficientNetV2 paper
+    def __init__(
+        self,
+        expand_ratio: float,
+        kernel: int,
+        stride: int,
+        input_channels: int,
+        out_channels: int,
+        num_layers: int,
+        width_mult: float = 1.0,
+        depth_mult: float = 1.0,
+        block: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        input_channels = self.adjust_channels(input_channels, width_mult)
+        out_channels = self.adjust_channels(out_channels, width_mult)
+        num_layers = self.adjust_depth(num_layers, depth_mult)
+        if block is None:
+            block = MBConv
+        super().__init__(
+            expand_ratio,
+            kernel,
+            stride,
+            input_channels,
+            out_channels,
+            num_layers,
+            block,
+        )
 
-class FrozenBatchNorm2d(torch.nn.Module):
-    """
-    BatchNorm2d where the batch statistics and the affine parameters are fixed
+    @staticmethod
+    def adjust_depth(num_layers: int, depth_mult: float):
+        return int(math.ceil(num_layers * depth_mult))
 
-    Args:
-        num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)``
-        eps (float): a value added to the denominator for numerical stability. Default: 1e-5
-    """
 
+class FusedMBConvConfig(_MBConvConfig):
+    # Stores information listed at Table 4 of the EfficientNetV2 paper
     def __init__(
         self,
-        num_features: int,
-        eps: float = 1e-5,
-    ):
-        super().__init__()
-        # _log_api_usage_once(self)
-        self.eps = eps
-        self.register_buffer("weight", torch.ones(num_features))
-        self.register_buffer("bias", torch.zeros(num_features))
-        self.register_buffer("running_mean", torch.zeros(num_features))
-        self.register_buffer("running_var", torch.ones(num_features))
-
-    def _load_from_state_dict(
-        self,
-        state_dict: dict,
-        prefix: str,
-        local_metadata: dict,
-        strict: bool,
-        missing_keys: List[str],
-        unexpected_keys: List[str],
-        error_msgs: List[str],
-    ):
-        num_batches_tracked_key = prefix + "num_batches_tracked"
-        if num_batches_tracked_key in state_dict:
-            del state_dict[num_batches_tracked_key]
-
-        super()._load_from_state_dict(
-            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        expand_ratio: float,
+        kernel: int,
+        stride: int,
+        input_channels: int,
+        out_channels: int,
+        num_layers: int,
+        block: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        if block is None:
+            block = FusedMBConv
+        super().__init__(
+            expand_ratio,
+            kernel,
+            stride,
+            input_channels,
+            out_channels,
+            num_layers,
+            block,
         )
 
-    def forward(self, x: Tensor) -> Tensor:
-        # move reshapes to the beginning
-        # to make it fuser-friendly
-        w = self.weight.reshape(1, -1, 1, 1)
-        b = self.bias.reshape(1, -1, 1, 1)
-        rv = self.running_var.reshape(1, -1, 1, 1)
-        rm = self.running_mean.reshape(1, -1, 1, 1)
-        scale = w * (rv + self.eps).rsqrt()
-        bias = b - rm * scale
-        return x * scale + bias
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})"
-
 
-class ConvNormActivation(torch.nn.Sequential):
+class MBConv(nn.Module):
     def __init__(
         self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int = 3,
-        stride: int = 1,
-        padding: Optional[int] = None,
-        groups: int = 1,
-        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
-        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
-        dilation: int = 1,
-        inplace: Optional[bool] = True,
-        bias: Optional[bool] = None,
-        conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d,
+        cnf: MBConvConfig,
+        stochastic_depth_prob: float,
+        norm_layer: Callable[..., nn.Module],
+        se_layer: Callable[..., nn.Module] = SqueezeExcitation,
     ) -> None:
+        super().__init__()
 
-        if padding is None:
-            padding = (kernel_size - 1) // 2 * dilation
-        if bias is None:
-            bias = norm_layer is None
+        if not (1 <= cnf.stride <= 2):
+            raise ValueError("illegal stride value")
 
-        layers = [
-            conv_layer(
-                in_channels,
-                out_channels,
-                kernel_size,
-                stride,
-                padding,
-                dilation=dilation,
-                groups=groups,
-                bias=bias,
-            )
-        ]
+        self.use_res_connect = (
+            cnf.stride == 1 and cnf.input_channels == cnf.out_channels
+        )
 
-        if norm_layer is not None:
-            layers.append(norm_layer(out_channels))
+        layers: List[nn.Module] = []
+        activation_layer = nn.SiLU
 
-        if activation_layer is not None:
-            params = {} if inplace is None else {"inplace": inplace}
-            layers.append(activation_layer(**params))
-        super().__init__(*layers)
-        # _log_api_usage_once(self)
-        self.out_channels = out_channels
+        # expand
+        expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio)
+        if expanded_channels != cnf.input_channels:
+            layers.append(
+                Conv2dNormActivation(
+                    cnf.input_channels,
+                    expanded_channels,
+                    kernel_size=1,
+                    norm_layer=norm_layer,
+                    activation_layer=activation_layer,
+                )
+            )
 
-        if self.__class__ == ConvNormActivation:
-            warnings.warn(
-                "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead."
+        # depthwise
+        layers.append(
+            Conv2dNormActivation(
+                expanded_channels,
+                expanded_channels,
+                kernel_size=cnf.kernel,
+                stride=cnf.stride,
+                groups=expanded_channels,
+                norm_layer=norm_layer,
+                activation_layer=activation_layer,
             )
+        )
 
+        # squeeze and excitation
+        squeeze_channels = max(1, cnf.input_channels // 4)
+        layers.append(
+            se_layer(
+                expanded_channels,
+                squeeze_channels,
+                activation=partial(nn.SiLU, inplace=True),
+            )
+        )
 
-class Conv2dNormActivation(ConvNormActivation):
-    """
-    Configurable block used for Convolution2d-Normalization-Activation blocks.
+        # project
+        layers.append(
+            Conv2dNormActivation(
+                expanded_channels,
+                cnf.out_channels,
+                kernel_size=1,
+                norm_layer=norm_layer,
+                activation_layer=None,
+            )
+        )
 
-    Args:
-        in_channels (int): Number of channels in the input image
-        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
-        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
-        stride (int, optional): Stride of the convolution. Default: 1
-        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
-        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
-        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm2d``
-        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
-        dilation (int): Spacing between kernel elements. Default: 1
-        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
-        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+        self.block = nn.Sequential(*layers)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+        self.out_channels = cnf.out_channels
 
-    """
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.block(input)
+        if self.use_res_connect:
+            result = self.stochastic_depth(result)
+            result += input
+        return result
 
+
+class FusedMBConv(nn.Module):
     def __init__(
         self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int = 3,
-        stride: int = 1,
-        padding: Optional[int] = None,
-        groups: int = 1,
-        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
-        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
-        dilation: int = 1,
-        inplace: Optional[bool] = True,
-        bias: Optional[bool] = None,
+        cnf: FusedMBConvConfig,
+        stochastic_depth_prob: float,
+        norm_layer: Callable[..., nn.Module],
     ) -> None:
+        super().__init__()
 
-        super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            groups,
-            norm_layer,
-            activation_layer,
-            dilation,
-            inplace,
-            bias,
-            torch.nn.Conv2d,
+        if not (1 <= cnf.stride <= 2):
+            raise ValueError("illegal stride value")
+
+        self.use_res_connect = (
+            cnf.stride == 1 and cnf.input_channels == cnf.out_channels
         )
 
+        layers: List[nn.Module] = []
+        activation_layer = nn.SiLU
 
-class Conv3dNormActivation(ConvNormActivation):
-    """
-    Configurable block used for Convolution3d-Normalization-Activation blocks.
+        expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio)
+        if expanded_channels != cnf.input_channels:
+            # fused expand
+            layers.append(
+                Conv2dNormActivation(
+                    cnf.input_channels,
+                    expanded_channels,
+                    kernel_size=cnf.kernel,
+                    stride=cnf.stride,
+                    norm_layer=norm_layer,
+                    activation_layer=activation_layer,
+                )
+            )
 
-    Args:
-        in_channels (int): Number of channels in the input video.
-        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
-        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
-        stride (int, optional): Stride of the convolution. Default: 1
-        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
-        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
-        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm3d``
-        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
-        dilation (int): Spacing between kernel elements. Default: 1
-        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
-        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
-    """
+            # project
+            layers.append(
+                Conv2dNormActivation(
+                    expanded_channels,
+                    cnf.out_channels,
+                    kernel_size=1,
+                    norm_layer=norm_layer,
+                    activation_layer=None,
+                )
+            )
+        else:
+            layers.append(
+                Conv2dNormActivation(
+                    cnf.input_channels,
+                    cnf.out_channels,
+                    kernel_size=cnf.kernel,
+                    stride=cnf.stride,
+                    norm_layer=norm_layer,
+                    activation_layer=activation_layer,
+                )
+            )
 
+        self.block = nn.Sequential(*layers)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+        self.out_channels = cnf.out_channels
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.block(input)
+        if self.use_res_connect:
+            result = self.stochastic_depth(result)
+            result += input
+        return result
+
+
+class EfficientNetConfig(PretrainedConfig):
+    
+    model_type = "efficientnet"
+    
     def __init__(
         self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int = 3,
-        stride: int = 1,
-        padding: Optional[int] = None,
-        groups: int = 1,
-        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm3d,
-        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
-        dilation: int = 1,
-        inplace: Optional[bool] = True,
-        bias: Optional[bool] = None,
+        # inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
+        dropout: float=0.25,
+        num_channels: int = 61,
+        stochastic_depth_prob: float = 0.2,
+        num_classes: int = 2,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        # last_channel: Optional[int] = None,
+        size: str='v2_s',
+        width_mult: float = 1.0,
+        depth_mult: float = 1.0,
+        **kwargs: Any,
     ) -> None:
+        """
+        EfficientNet V1 and V2 main class
+
+        Args:
+            inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure
+            dropout (float): The droupout probability
+            stochastic_depth_prob (float): The stochastic depth probability
+            num_classes (int): Number of classes
+            norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use
+            last_channel (int): The number of channels on the penultimate layer
+        """
+        
+        
+        # self.model = EfficientNet(
+        #                             dropout=dropout,
+        #                             num_channels=num_channels,
+        #                             num_classes=num_classes,
+        #                             size=size,
+        #                             stochastic_depth_prob=stochastic_depth_prob,
+        #                             width_mult=width_mult,
+        #                             depth_mult=depth_mult,
+        # )
+        
+        # 
+        self.dropout=dropout
+        self.num_channels=num_channels
+        self.num_classes=num_classes
+        self.size=size
+        self.stochastic_depth_prob=stochastic_depth_prob
+        self.width_mult=width_mult
+        self.depth_mult=depth_mult
+                                    
+        super().__init__(**kwargs)
+        
+    
+class EfficientNetPreTrained(PreTrainedModel):
+    
+    config_class = EfficientNetConfig
+    
+    def __init__(
+        self,
+        config
+    ):
+        super().__init__(config)   
+        self.model = EfficientNet(  dropout=config.dropout,
+                                    num_channels=config.num_channels,
+                                    num_classes=config.num_classes,
+                                    size=config.size,
+                                    stochastic_depth_prob=config.stochastic_depth_prob,
+                                    width_mult=config.width_mult,
+                                    depth_mult=config.depth_mult,)
+        
+    def forward(self, tensor):
+        return self.model.forward(tensor)
+    
+    
+class EfficientNet(nn.Module):
+    
+    
+    def __init__(
+        self,
+        # inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
+        dropout: float=0.25,
+        num_channels: int = 61,
+        stochastic_depth_prob: float = 0.2,
+        num_classes: int = 2,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        # last_channel: Optional[int] = None,
+        size: str='v2_s',
+        width_mult: float = 1.0,
+        depth_mult: float = 1.0,
+        **kwargs: Any,
+    ) -> None:
+        """
+        EfficientNet V1 and V2 main class
+
+        Args:
+            inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure
+            dropout (float): The droupout probability
+            stochastic_depth_prob (float): The stochastic depth probability
+            num_classes (int): Number of classes
+            norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use
+            last_channel (int): The number of channels on the penultimate layer
+        """
+        super().__init__()
+        # _log_api_usage_once(self)
+        
+        inverted_residual_setting, last_channel = _efficientnet_conf(
+                     "efficientnet_%s" % (size), width_mult=width_mult, depth_mult=depth_mult
+                    )
+
+        if not inverted_residual_setting:
+            raise ValueError("The inverted_residual_setting should not be empty")
+        elif not (
+            isinstance(inverted_residual_setting, Sequence)
+            and all([isinstance(s, _MBConvConfig) for s in inverted_residual_setting])
+        ):
+            raise TypeError(
+                "The inverted_residual_setting should be List[MBConvConfig]"
+            )
+
+        if "block" in kwargs:
+            warnings.warn(
+                "The parameter 'block' is deprecated since 0.13 and will be removed 0.15. "
+                "Please pass this information on 'MBConvConfig.block' instead."
+            )
+            if kwargs["block"] is not None:
+                for s in inverted_residual_setting:
+                    if isinstance(s, MBConvConfig):
+                        s.block = kwargs["block"]
 
-        super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            groups,
-            norm_layer,
-            activation_layer,
-            dilation,
-            inplace,
-            bias,
-            torch.nn.Conv3d,
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        layers: List[nn.Module] = []
+
+        # building first layer
+        firstconv_output_channels = inverted_residual_setting[0].input_channels
+        layers.append(
+            Conv2dNormActivation(
+                num_channels,
+                firstconv_output_channels,
+                kernel_size=3,
+                stride=2,
+                norm_layer=norm_layer,
+                activation_layer=nn.SiLU,
+            )
         )
 
+        # building inverted residual blocks
+        total_stage_blocks = sum(cnf.num_layers for cnf in inverted_residual_setting)
+        stage_block_id = 0
+        for cnf in inverted_residual_setting:
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # copy to avoid modifications. shallow copy is enough
+                block_cnf = copy.copy(cnf)
 
-class SqueezeExcitation(torch.nn.Module):
-    """
-    This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1).
-    Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in eq. 3.
+                # overwrite info if not the first conv in the stage
+                if stage:
+                    block_cnf.input_channels = block_cnf.out_channels
+                    block_cnf.stride = 1
 
-    Args:
-        input_channels (int): Number of channels in the input image
-        squeeze_channels (int): Number of squeeze channels
-        activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU``
-        scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid``
-    """
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = (
+                    stochastic_depth_prob * float(stage_block_id) / total_stage_blocks
+                )
 
-    def __init__(
-        self,
-        input_channels: int,
-        squeeze_channels: int,
-        activation: Callable[..., torch.nn.Module] = torch.nn.ReLU,
-        scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid,
-    ) -> None:
-        super().__init__()
-        # _log_api_usage_once(self)
-        self.avgpool = torch.nn.AdaptiveAvgPool2d(1)
-        self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1)
-        self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1)
-        self.activation = activation()
-        self.scale_activation = scale_activation()
+                stage.append(block_cnf.block(block_cnf, sd_prob, norm_layer))
+                stage_block_id += 1
 
-    def _scale(self, input: Tensor) -> Tensor:
-        scale = self.avgpool(input)
-        scale = self.fc1(scale)
-        scale = self.activation(scale)
-        scale = self.fc2(scale)
-        return self.scale_activation(scale)
+            layers.append(nn.Sequential(*stage))
 
-    def forward(self, input: Tensor) -> Tensor:
-        scale = self._scale(input)
-        return scale * input
+        # building last several layers
+        lastconv_input_channels = inverted_residual_setting[-1].out_channels
+        lastconv_output_channels = (
+            last_channel if last_channel is not None else 4 * lastconv_input_channels
+        )
+        layers.append(
+            Conv2dNormActivation(
+                lastconv_input_channels,
+                lastconv_output_channels,
+                kernel_size=1,
+                norm_layer=norm_layer,
+                activation_layer=nn.SiLU,
+            )
+        )
 
+        self.features = nn.Sequential(*layers)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout, inplace=True),
+            nn.Linear(lastconv_output_channels, num_classes),
+        )
 
-class MLP(torch.nn.Sequential):
-    """This block implements the multi-layer perceptron (MLP) module.
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out")
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                init_range = 1.0 / math.sqrt(m.out_features)
+                nn.init.uniform_(m.weight, -init_range, init_range)
+                nn.init.zeros_(m.bias)
+                
+        # super().__init__(**kwargs)
 
-    Args:
-        in_channels (int): Number of channels of the input
-        hidden_channels (List[int]): List of the hidden channel dimensions
-        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``None``
-        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
-        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
-        bias (bool): Whether to use bias in the linear layer. Default ``True``
-        dropout (float): The probability for the dropout layer. Default: 0.0
-    """
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.features(x)
 
-    def __init__(
-        self,
-        in_channels: int,
-        hidden_channels: List[int],
-        norm_layer: Optional[Callable[..., torch.nn.Module]] = None,
-        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
-        inplace: Optional[bool] = True,
-        bias: bool = True,
-        dropout: float = 0.0,
-    ):
-        # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal:
-        # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py
-        params = {} if inplace is None else {"inplace": inplace}
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
 
-        layers = []
-        in_dim = in_channels
-        for hidden_dim in hidden_channels[:-1]:
-            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
-            if norm_layer is not None:
-                layers.append(norm_layer(hidden_dim))
-            layers.append(activation_layer(**params))
-            layers.append(torch.nn.Dropout(dropout, **params))
-            in_dim = hidden_dim
+        x = self.classifier(x)
 
-        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias))
-        layers.append(torch.nn.Dropout(dropout, **params))
+        return x
 
-        super().__init__(*layers)
-        # _log_api_usage_once(self)
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
 
 
-class Permute(torch.nn.Module):
-    """This module returns a view of the tensor input with its dimensions permuted.
+# def _efficientnet(
+#     inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]],
+#     dropout: float,
+#     last_channel: Optional[int],
+#     weights=None,
+#     num_channels: int = 61,
+#     stochastic_depth_prob: float = 0.2,
+#     progress: bool = True,
+#     num_classes: int = 2,
+#     **kwargs: Any,
+# ) -> EfficientNetCongig:
 
-    Args:
-        dims (List[int]): The desired ordering of dimensions
-    """
+#     model = EfficientNetCongif(
+#         inverted_residual_setting,
+#         dropout,
+#         num_classes=num_classes,
+#         num_channels=num_channels,
+#         stochastic_depth_prob=stochastic_depth_prob,
+#         last_channel=last_channel,
+#         **kwargs,
+#     )
 
-    def __init__(self, dims: List[int]):
-        super().__init__()
-        self.dims = dims
+#     return model
 
-    def forward(self, x: Tensor) -> Tensor:
-        return torch.permute(x, self.dims)
 
+def _efficientnet_conf(
+    arch: str,
+    **kwargs: Any,
+) -> Tuple[Sequence[Union[MBConvConfig, FusedMBConvConfig]], Optional[int]]:
+    inverted_residual_setting: Sequence[Union[MBConvConfig, FusedMBConvConfig]]
+    if arch.startswith("efficientnet_b"):
+        bneck_conf = partial(
+            MBConvConfig,
+            width_mult=kwargs.pop("width_mult"),
+            depth_mult=kwargs.pop("depth_mult"),
+        )
+        inverted_residual_setting = [
+            bneck_conf(1, 3, 1, 32, 16, 1),
+            bneck_conf(6, 3, 2, 16, 24, 2),
+            bneck_conf(6, 5, 2, 24, 40, 2),
+            bneck_conf(6, 3, 2, 40, 80, 3),
+            bneck_conf(6, 5, 1, 80, 112, 3),
+            bneck_conf(6, 5, 2, 112, 192, 4),
+            bneck_conf(6, 3, 1, 192, 320, 1),
+        ]
+        last_channel = None
+    elif arch.startswith("efficientnet_v2_s"):
+        inverted_residual_setting = [
+            FusedMBConvConfig(1, 3, 1, 24, 24, 2),
+            FusedMBConvConfig(4, 3, 2, 24, 48, 4),
+            FusedMBConvConfig(4, 3, 2, 48, 64, 4),
+            MBConvConfig(4, 3, 2, 64, 128, 6),
+            MBConvConfig(6, 3, 1, 128, 160, 9),
+            MBConvConfig(6, 3, 2, 160, 256, 15),
+        ]
+        last_channel = 1280
+    elif arch.startswith("efficientnet_v2_m"):
+        inverted_residual_setting = [
+            FusedMBConvConfig(1, 3, 1, 24, 24, 3),
+            FusedMBConvConfig(4, 3, 2, 24, 48, 5),
+            FusedMBConvConfig(4, 3, 2, 48, 80, 5),
+            MBConvConfig(4, 3, 2, 80, 160, 7),
+            MBConvConfig(6, 3, 1, 160, 176, 14),
+            MBConvConfig(6, 3, 2, 176, 304, 18),
+            MBConvConfig(6, 3, 1, 304, 512, 5),
+        ]
+        last_channel = 1280
+    elif arch.startswith("efficientnet_v2_l"):
+        inverted_residual_setting = [
+            FusedMBConvConfig(1, 3, 1, 32, 32, 4),
+            FusedMBConvConfig(4, 3, 2, 32, 64, 7),
+            FusedMBConvConfig(4, 3, 2, 64, 96, 7),
+            MBConvConfig(4, 3, 2, 96, 192, 10),
+            MBConvConfig(6, 3, 1, 192, 224, 19),
+            MBConvConfig(6, 3, 2, 224, 384, 25),
+            MBConvConfig(6, 3, 1, 384, 640, 7),
+        ]
+        last_channel = 1280
+    else:
+        raise ValueError(f"Unsupported model type {arch}")
+
+    return inverted_residual_setting, last_channel
 
 
 
+##### normal stuff ####
 
 def normalize_array(x: list):