add toggle for using skip connections

models/convnext.py

@@ -180,9 +180,12 @@ class ConvNeXtReverseBase2(nn.Module):
         block: Optional[Callable[..., nn.Module]] = None,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
         device: str = 'cuda',
+        use_skip=False,
     ) -> None:
         super().__init__()
 
+        self.use_skip = use_skip
+
         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])):
@@ -259,10 +262,13 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x, features = x
-        x = self.first(x)
+        input, features = x
+        x = self.first(input)
         for i in range(5 - 2):
-            x = self.stages[i](x + features[2-i])
+            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)

models/depth_estim.py

@@ -16,18 +16,20 @@ class ViewEncoder(nn.Module):
         self.lin1 = nn.LazyLinear(features).to(device)
         self.act1 = nn.GELU().to(device)
         self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
+        self.act2 = nn.Tanh().to(device)
     
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         x: torch.Tensor = self.lin1(input)
         x = self.act1(x)
         x = self.lin2(x)
+        x = self.act2(x)
         x = x.view(x.shape[0],self.features,15,20)
 
         return x
 
 class DepthEstimation(nn.Module):
-    def __init__(self, view_features=128, device='cuda'):
+    def __init__(self, view_features=128, device='cuda', use_skip=False):
         super().__init__()
 
         
@@ -40,16 +42,16 @@ class DepthEstimation(nn.Module):
         
         block_setting_reverse = [
             CNBlockConfig(768 + view_features, 384, 3),
-            CNBlockConfig(384, 192, 9),
-            CNBlockConfig(192, 96, 3),
-            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(384 * ( 2 if use_skip else 1), 192, 9),
+            CNBlockConfig(192 * ( 2 if use_skip else 1), 96, 3),
+            CNBlockConfig(96 * ( 2 if use_skip else 1), 27, 3),
             CNBlockConfig(27, 9, 3),
             CNBlockConfig(9, None, 3),
         ]
 
         self.convnext = ConvNeXt(block_setting, device=device)
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device)
+        self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)
 
         self.view_encoder = ViewEncoder(view_features)