trying skip connection

models/convnext_orig.py

@@ -81,8 +81,6 @@ class ConvNeXtBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        layers_list: List[nn.Module] = []
-
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
@@ -95,7 +93,8 @@ class ConvNeXtBase(nn.Module):
                 activation_layer=None,
                 bias=True,
             ).to(device)
-        layers_list.append(self.stem)
+        self.stages = []
+        self.downsamples = []
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
@@ -107,17 +106,20 @@ 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
-            layers_list.append(nn.Sequential(*stage).to(device))
+            self.stages.append(nn.Sequential(*stage).to(device))
             if cnf.out_channels is not None:
                 # Downsampling
-                layers_list.append(
+                self.downsamples.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                     ).to(device)
                 )
+        self.stages = nn.ModuleList(self.stages)
+        self.downsamples = nn.ModuleList(self.downsamples)
 
-        self.layers = nn.Sequential(*layers_list).to(device)
+        # self.layers = nn.Sequential(*layers_list).to(device)
+        # self.layers = nn.ModuleList(layers_list)
         self.act = nn.Tanh().to(device)
 
         # for m in self.modules():
@@ -127,9 +129,16 @@ class ConvNeXtBase(nn.Module):
         #             nn.init.zeros_(m.bias)
 
     def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.layers(x)
+        x = self.stem(x)
+        features = []
+        for i in range(len(self.stages)):
+            x = self.stages[i](x)
+            if i == len(self.stages) - 1:
+                continue
+            features.append(x)
+            x = self.downsamples[i](x)
         x = self.act(x)
-        return x
+        return x, features
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
@@ -165,13 +174,14 @@ class ConvNeXtReverseBase(nn.Module):
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
         self.stem = nn.Sequential(
-                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=4, stride=4),
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
                         norm_layer(block_setting[0].out_channels),
                     ).to(device)
         
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        layers_list: List[nn.Module] = []
+        self.stages = []
+        self.upsamples = []
 
 
         for i, cnf in enumerate(block_setting):
@@ -182,17 +192,17 @@ class ConvNeXtReverseBase(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
-            layers_list.append(nn.Sequential(*stage))
+            self.stages.append(nn.Sequential(*stage).to(device))
 
             if i == 0:
-                layers_list.append(self.stem)
+                self.upsamples.append(self.stem)
             else:
                 # Upsampling
-                layers_list.append(
+                self.upsamples.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    )
+                    ).to(device)
                 )
 
         self.last = nn.Sequential(
@@ -201,7 +211,8 @@ class ConvNeXtReverseBase(nn.Module):
                         Permute([0, 3, 1, 2]),
                     ).to(device)
         
-        self.layers = nn.Sequential(*layers_list).to(device)
+        self.stages = nn.ModuleList(self.stages)
+        self.upsamples = nn.ModuleList(self.upsamples)
         self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device)
 
 
@@ -213,7 +224,15 @@ class ConvNeXtReverseBase(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x = self.layers(x)
+        x, features = x
+        for i in range(len(self.stages)-1):
+            if i == 0:
+                x = self.stages[i](x)
+            else:
+                x = self.stages[i](torch.cat([x, features[-i]], dim=1))
+            x = self.upsamples[i](x)
+        x = self.stages[-1](x)
+        x = self.upsamples[-1](x)
         x = self.last(x)
         x = self.activation(x)
         return x