fix convenxt layers

models/convnext.py

@@ -8,25 +8,10 @@ from torch import nn, Tensor
 from torch.nn import functional as F
 
 from torchvision.ops.misc import Conv2dNormActivation, Permute
-# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
-# from torchvision.ops import Permute
 from torchvision.ops.stochastic_depth import StochasticDepth
 from utils.convnext_utils import CNBlockConfig
 
 
-__all__ = [
-    "ConvNeXt",
-    "ConvNeXt_Tiny_Weights",
-    "ConvNeXt_Small_Weights",
-    "ConvNeXt_Base_Weights",
-    "ConvNeXt_Large_Weights",
-    "convnext_tiny",
-    "convnext_small",
-    "convnext_base",
-    "convnext_large",
-]
-
-
 class LayerNorm2d(nn.LayerNorm):
     def forward(self, x: Tensor) -> Tensor:
         x = x.permute(0, 2, 3, 1)
@@ -91,11 +76,7 @@ class ConvNeXtBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        # self.permute_cf = Permute([0, 3, 1, 2])
+        layers_list: List[nn.Module] = []
-        # self.permute_cl = Permute([0, 2, 3, 1])
-
-        self.stages: List[nn.Module] = []
-        self.downsamples: List[nn.Module] = []
 
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
@@ -121,17 +102,17 @@ class ConvNeXtBase(nn.Module):
                 sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                 stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
                 stage_block_id += 1
-            self.stages.append(nn.Sequential(*stage).to(device))
+            layers_list.append(nn.Sequential(*stage))
             if cnf.out_channels is not None:
                 # Downsampling
-                self.downsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    ).to(device)
+                    )
                 )
 
-        self.act = nn.Tanh()
+        self.layers = nn.Sequential(*layers_list).to(device)
 
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -140,25 +121,15 @@ class ConvNeXtBase(nn.Module):
                     nn.init.zeros_(m.bias)
 
     def _forward_impl(self, x: Tensor) -> Tensor:
-        # x = self.permute_cf(x)
         x = self.stem(x)
-        # features: List[Tensor] = []
+        x = self.layers(x)
-        for i in range(len(self.stages) - 1):
-            x = self.stages[i](x)
-            # features.append(x)
-            x = self.downsamples[i](x)
-        x = self.stages[-1](x)
-        x = self.act(x)
-
-        # x = self.permute_cl(x)
         return x
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
-
+class ConvNeXtReverseBase(nn.Module):
-class ConvNeXtReverseBase2(nn.Module):
     def __init__(
         self,
         block_setting: List[CNBlockConfig],
@@ -187,21 +158,6 @@ class ConvNeXtReverseBase2(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        # self.permute_cf = Permute([0, 3, 1, 2])
-        # self.permute_cl = Permute([0, 2, 3, 1])
-
-        # Stem
-        # firstconv_output_channels = block_setting[0].input_channels
-        # self.first = nn.Sequential(
-        #             nn.ConvTranspose2d(
-        #                 firstconv_output_channels,
-        #                 firstconv_output_channels,
-        #                 kernel_size=2,
-        #                 stride=2,
-        #                 padding=0,
-        #                 ),
-        #             norm_layer(firstconv_output_channels),
-        #         ).to(device)
         self.first = nn.Sequential(
                         nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
                         norm_layer(block_setting[0].out_channels),
@@ -210,8 +166,8 @@ class ConvNeXtReverseBase2(nn.Module):
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        self.stages: List[nn.Module] = []
+        layers_list: List[nn.Module] = []
-        self.upsamples: List[nn.Module] = []
+        
         for cnf in block_setting[1:]:
             if cnf.input_channels is not None:
                 # Bottlenecks
@@ -221,18 +177,20 @@ class ConvNeXtReverseBase2(nn.Module):
                     sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                     stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
                     stage_block_id += 1
-                self.stages.append(nn.Sequential(*stage).to(device))
+                layers_list.append(nn.Sequential(*stage))
             else:
-                self.stages.append(nn.Identity().to(device))
+                layers_list.append(nn.Identity())
 
             if cnf.out_channels is not None:
                 # Upsampling
-                self.upsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    ).to(device)
                     )
+                )
+
+        self.layers = nn.Sequential(*layers_list).to(device)
 
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
@@ -240,13 +198,8 @@ class ConvNeXtReverseBase2(nn.Module):
                     kernel_size=3,
                     padding='same',
                     ).to(device)
-        self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
         
-        # self.last = nn.Sequential(
+        self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
-        #                 Permute([0, 2, 3, 1]),
-        #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
-        #                 Permute([0, 3, 1, 2]),
-        #             ).to(device)
 
 
         for m in self.modules():
@@ -257,134 +210,16 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        # input, features = x
-        # x = self.permute_cf(x)
-
         x = self.first(x)
-        for i in range(len(self.stages)):
+        x = self.layers(x)
-            # if self.use_skip:
-            #     x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
-            # else:
-            x = self.stages[i](x)
-            x = self.upsamples[i](x)
-        
-        # x = self.stages[-2](x)
-        # x = self.upsamples[-1](x)
-
-        # x = self.stages[-1](x)
         x = self.last(x)
         x = self.activation(x)
-        # x = self.permute_cl(x)
-        
         return x
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)
 
 
-class ConvNeXtReverseBase(nn.Module):
-    def __init__(
-        self,
-        block_setting: List[CNBlockConfig],
-        stochastic_depth_prob: float = 0.0,
-        layer_scale: float = 1e-6,
-        block: Optional[Callable[..., nn.Module]] = None,
-        norm_layer: Optional[Callable[..., nn.Module]] = None
-    ) -> None:
-        super().__init__()
-
-        if not block_setting:
-            raise ValueError("The block_setting should not be empty")
-        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
-            raise TypeError("The block_setting should be List[CNBlockConfig]")
-
-        if block is None:
-            block = CNBlock
-
-        if norm_layer is None:
-            norm_layer = partial(LayerNorm2d, eps=1e-6)
-
-        layers: List[nn.Module] = []
-
-
-        # add bottleneck which will be executed at the end, to process after the final larger upsampling
-
-
-
-        # Stem
-        firstconv_output_channels = block_setting[0].input_channels
-
-        layers.append(
-            nn.Sequential(
-                nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
-
-            )
-        )
-
-        stage: List[nn.Module] = []
-        for _ in range(3):
-            stage.append(block(16, layer_scale, 0.0))
-        layers.append(nn.Sequential(*stage))
-        
-        layers.append(
-            # Convtr NormActivation(
-            #     firstconv_output_channels,
-            #     3,
-            #     kernel_size=4,
-            #     stride=4,
-            #     padding=0,
-            #     norm_layer=None,
-            #     activation_layer=None,
-            #     bias=True,
-            # )
-            nn.Sequential(
-                nn.ConvTranspose2d(
-                    firstconv_output_channels,
-                    16,
-                    kernel_size=4,
-                    stride=4,
-                    padding=0,
-                    ),
-            )
-        )
-
-        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
-        stage_block_id = 0
-        for cnf in block_setting:
-            # Bottlenecks
-            stage: List[nn.Module] = []
-            for _ in range(cnf.num_layers):
-                # adjust stochastic depth probability based on the depth of the stage block
-                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
-                stage.append(block(cnf.input_channels, layer_scale, sd_prob))
-                stage_block_id += 1
-            layers.append(nn.Sequential(*stage))
-            if cnf.out_channels is not None:
-                # Upsampling
-                layers.append(
-                    nn.Sequential(
-                        norm_layer(cnf.out_channels),
-                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
-                    )
-                )
-
-        self.features = nn.Sequential(*reversed(layers))
-
-        for m in self.modules():
-            if isinstance(m, (nn.Conv2d, nn.Linear)):
-                nn.init.trunc_normal_(m.weight, std=0.02)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        # x = torch.tanh(x)
-        return x
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self._forward_impl(x)
-
-
 def _convnext(
     block_setting: List[CNBlockConfig],
     stochastic_depth_prob: float,
@@ -397,7 +232,7 @@ def _convnext(
 ) -> ConvNeXtBase:
 
     if reversed:
-        model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
     else:
         model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)