From 191f6b01e09ea81b0ef42f76124c80498b2ebd31 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 14 Jan 2023 17:13:12 +0000 Subject: [PATCH] add toggle for using skip connections --- models/convnext.py | 12 +++++++++--- models/depth_estim.py | 12 +++++++----- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/models/convnext.py b/models/convnext.py index 966926d..d18bb2c 100644 --- a/models/convnext.py +++ b/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) diff --git a/models/depth_estim.py b/models/depth_estim.py index 0ee0335..29961c7 100644 --- a/models/depth_estim.py +++ b/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)