add toggle for using skip connections

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)