From bca2e0bd2cdc905e55b393c57b40d1be0263529d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 9 Jan 2023 21:22:15 +0000 Subject: [PATCH] new reversed convnext --- data/dataset.py | 2 - models/convnext.py | 167 +++++++++++++++++++++++++++++++++++------- models/depth_estim.py | 27 +++++-- 3 files changed, 161 insertions(+), 35 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index 8f4b8ed..2592e18 100644 --- a/data/dataset.py +++ b/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, diff --git a/models/convnext.py b/models/convnext.py index 4d7b408..b24ccaa 100644 --- a/models/convnext.py +++ b/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: diff --git a/models/depth_estim.py b/models/depth_estim.py index 92cdcd2..281313f 100644 --- a/models/depth_estim.py +++ b/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