new reversed convnext

data/dataset.py

@@ -190,10 +190,8 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # ex.imshow(v2i_coverage).show()
 
-
         rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
 
-
         # image = self.transforms(image)
         data = { 'mesh_posed': mesh_posed,
                  'depth_image': depth_image, 

models/convnext.py

@@ -93,7 +93,8 @@ class ConvNeXtBase(nn.Module):
         stochastic_depth_prob: float = 0.0,
         layer_scale: float = 1e-6,
         block: Optional[Callable[..., nn.Module]] = None,
-        norm_layer: Optional[Callable[..., nn.Module]] = None
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
     ) -> None:
         super().__init__()
 
@@ -108,12 +109,12 @@ class ConvNeXtBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        layers: List[nn.Module] = []
+        self.stages: List[nn.Module] = []
+        self.downsamples: List[nn.Module] = []
 
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
-        layers.append(
-            Conv2dNormActivation(
+        self.stem = Conv2dNormActivation(
                 3,
                 firstconv_output_channels,
                 kernel_size=4,
@@ -122,8 +123,8 @@ class ConvNeXtBase(nn.Module):
                 norm_layer=norm_layer,
                 activation_layer=None,
                 bias=True,
-            )
-        )
+            ).to(device)
+        
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
@@ -135,18 +136,16 @@ 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))
                 stage_block_id += 1
-            layers.append(nn.Sequential(*stage))
+            self.stages.append(nn.Sequential(*stage).to(device))
             if cnf.out_channels is not None:
                 # Downsampling
-                layers.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.features = nn.Sequential(*layers)
-
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
                 nn.init.trunc_normal_(m.weight, std=0.02)
@@ -154,14 +153,123 @@ class ConvNeXtBase(nn.Module):
                     nn.init.zeros_(m.bias)
 
     def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        return x
+        x = self.stem(x)
+        features: List[Tensor] = []
+        for i in range(4 - 1):
+            x = self.stages[i](x)
+            features.append(x)
+            x = self.downsamples[i](x)
+        x = self.stages[-1](x)
+        
+
+        return x, features
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
 
+class ConvNeXtReverseBase2(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,
+        device: str = 'cuda',
+    ) -> 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)
+
+        # 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),
+                    ).to(device)
+
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        self.stages: List[nn.Module] = []
+        self.upsamples: List[nn.Module] = []
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # 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
+                self.stages.append(nn.Sequential(*stage).to(device))
+            else:
+                self.stages.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upnsampling
+                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.Conv2d(
+                    block_setting[-1].input_channels,
+                    1,
+                    kernel_size=7,
+                    padding=3,
+                    ).to(device)
+
+                    # Dense?
+
+
+        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:
+        x, features = x
+        x = self.first(x)
+        for i in range(5 - 2):
+            x = self.stages[i](x + features[2-i])
+            x = self.upsamples[i](x)
+        
+        x = self.stages[-2](x)
+        x = self.upsamples[-1](x)
+        x = self.stages[-1](x)
+        x = self.last(x)
+        return x
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
 class ConvNeXtReverseBase(nn.Module):
     def __init__(
         self,
@@ -188,14 +296,24 @@ class ConvNeXtReverseBase(nn.Module):
 
 
         # add bottleneck which will be executed at the end, to process after the final larger upsampling
-        stage: List[nn.Module] = []
-        for _ in range(3):
-            stage.append(block(1, layer_scale, 0.0))
-        layers.append(nn.Sequential(*stage))
+
 
 
         # 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,
@@ -210,7 +328,7 @@ class ConvNeXtReverseBase(nn.Module):
             nn.Sequential(
                 nn.ConvTranspose2d(
                     firstconv_output_channels,
-                    1,
+                    16,
                     kernel_size=4,
                     stride=4,
                     padding=0,
@@ -248,7 +366,7 @@ class ConvNeXtReverseBase(nn.Module):
 
     def _forward_impl(self, x: Tensor) -> Tensor:
         x = self.features(x)
-        x = torch.tanh(x)
+        # x = torch.tanh(x)
         return x
 
     def forward(self, x: Tensor) -> Tensor:
@@ -264,7 +382,7 @@ def _convnext(
 ) -> ConvNeXtBase:
 
     if reversed:
-        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+        model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
     else:
         model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
 
@@ -275,16 +393,9 @@ def _convnext(
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
         super().__init__()
 
-
-        block_setting = [
-            CNBlockConfig(96, 192, 3),
-            CNBlockConfig(192, 384, 3),
-            CNBlockConfig(384, 768, 9),
-            CNBlockConfig(768, None, 3),
-        ]
         self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:

models/depth_estim.py

@@ -6,21 +6,38 @@ import torch.nn as nn
 from torch.nn import functional as F
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from models.convnext import ConvNeXt
+from models.convnext import ConvNeXt, CNBlockConfig
 
 
 class DepthEstimation(nn.Module):
     def __init__(self, device='cuda'):
         super().__init__()
 
-        self.convnext = ConvNeXt()
+        
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, 768, 9),
+            CNBlockConfig(768, None, 3),
+        ]
+        
+        block_setting_reverse = [
+            CNBlockConfig(768, 384, 3),
+            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(27, 9, 3),
+            CNBlockConfig(9, None, 3),
+        ]
+
+        self.convnext = ConvNeXt(block_setting)
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(reversed=True)
+        self.decoder = ConvNeXt(block_setting_reverse, reversed=True)
 
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.convnext(x)
+        x, features = self.convnext(x)
         # x = self.flatten(x)
-        x = self.decoder(x)
+        x = self.decoder([x, features])
         
         return x