diff --git a/models/convnext.py b/models/convnext.py index e055370..b126f4c 100644 --- a/models/convnext.py +++ b/models/convnext.py @@ -8,25 +8,10 @@ from torch import nn, Tensor from torch.nn import functional as F from torchvision.ops.misc import Conv2dNormActivation, Permute -# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation -# from torchvision.ops import Permute from torchvision.ops.stochastic_depth import StochasticDepth from utils.convnext_utils import CNBlockConfig -__all__ = [ - "ConvNeXt", - "ConvNeXt_Tiny_Weights", - "ConvNeXt_Small_Weights", - "ConvNeXt_Base_Weights", - "ConvNeXt_Large_Weights", - "convnext_tiny", - "convnext_small", - "convnext_base", - "convnext_large", -] - - class LayerNorm2d(nn.LayerNorm): def forward(self, x: Tensor) -> Tensor: x = x.permute(0, 2, 3, 1) @@ -91,11 +76,7 @@ class ConvNeXtBase(nn.Module): if norm_layer is None: norm_layer = partial(LayerNorm2d, eps=1e-6) - # self.permute_cf = Permute([0, 3, 1, 2]) - # self.permute_cl = Permute([0, 2, 3, 1]) - - self.stages: List[nn.Module] = [] - self.downsamples: List[nn.Module] = [] + layers_list: List[nn.Module] = [] # Stem firstconv_output_channels = block_setting[0].input_channels @@ -121,17 +102,17 @@ 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, kernel_size)) stage_block_id += 1 - self.stages.append(nn.Sequential(*stage).to(device)) + layers_list.append(nn.Sequential(*stage)) if cnf.out_channels is not None: # Downsampling - self.downsamples.append( + layers_list.append( nn.Sequential( norm_layer(cnf.input_channels), nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), - ).to(device) + ) ) - self.act = nn.Tanh() + self.layers = nn.Sequential(*layers_list).to(device) for m in self.modules(): if isinstance(m, (nn.Conv2d, nn.Linear)): @@ -140,25 +121,15 @@ class ConvNeXtBase(nn.Module): nn.init.zeros_(m.bias) def _forward_impl(self, x: Tensor) -> Tensor: - # x = self.permute_cf(x) x = self.stem(x) - # features: List[Tensor] = [] - for i in range(len(self.stages) - 1): - x = self.stages[i](x) - # features.append(x) - x = self.downsamples[i](x) - x = self.stages[-1](x) - x = self.act(x) - - # x = self.permute_cl(x) + x = self.layers(x) return x def forward(self, x: Tensor) -> Tensor: return self._forward_impl(x) - -class ConvNeXtReverseBase2(nn.Module): +class ConvNeXtReverseBase(nn.Module): def __init__( self, block_setting: List[CNBlockConfig], @@ -187,21 +158,6 @@ class ConvNeXtReverseBase2(nn.Module): if norm_layer is None: norm_layer = partial(LayerNorm2d, eps=1e-6) - # self.permute_cf = Permute([0, 3, 1, 2]) - # self.permute_cl = Permute([0, 2, 3, 1]) - - # 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), @@ -210,8 +166,8 @@ class ConvNeXtReverseBase2(nn.Module): 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] = [] + layers_list: List[nn.Module] = [] + for cnf in block_setting[1:]: if cnf.input_channels is not None: # Bottlenecks @@ -221,33 +177,30 @@ class ConvNeXtReverseBase2(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, kernel_size)) stage_block_id += 1 - self.stages.append(nn.Sequential(*stage).to(device)) + layers_list.append(nn.Sequential(*stage)) else: - self.stages.append(nn.Identity().to(device)) + layers_list.append(nn.Identity()) if cnf.out_channels is not None: # Upsampling - self.upsamples.append( + layers_list.append( nn.Sequential( norm_layer(cnf.input_channels), nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), - ).to(device) + ) ) + self.layers = nn.Sequential(*layers_list).to(device) + self.last = nn.Conv2d( block_setting[-1].input_channels, 3 if not is_label else nr_of_labels, kernel_size=3, padding='same', ).to(device) + self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1) - # self.last = nn.Sequential( - # Permute([0, 2, 3, 1]), - # nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True), - # Permute([0, 3, 1, 2]), - # ).to(device) - for m in self.modules(): if isinstance(m, (nn.Conv2d, nn.Linear)): @@ -257,134 +210,16 @@ class ConvNeXtReverseBase2(nn.Module): def _forward_impl(self, x) -> Tensor: - # input, features = x - # x = self.permute_cf(x) - x = self.first(x) - for i in range(len(self.stages)): - # 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) - # x = self.upsamples[-1](x) - - # x = self.stages[-1](x) + x = self.layers(x) x = self.last(x) x = self.activation(x) - # x = self.permute_cl(x) - return x def forward(self, x) -> Tensor: return self._forward_impl(x) -class ConvNeXtReverseBase(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 - ) -> 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) - - layers: List[nn.Module] = [] - - - # add bottleneck which will be executed at the end, to process after the final larger upsampling - - - - # 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, - # 3, - # kernel_size=4, - # stride=4, - # padding=0, - # norm_layer=None, - # activation_layer=None, - # bias=True, - # ) - nn.Sequential( - nn.ConvTranspose2d( - firstconv_output_channels, - 16, - kernel_size=4, - stride=4, - padding=0, - ), - ) - ) - - total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) - stage_block_id = 0 - for cnf in block_setting: - # 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 - layers.append(nn.Sequential(*stage)) - if cnf.out_channels is not None: - # Upsampling - layers.append( - nn.Sequential( - norm_layer(cnf.out_channels), - nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2), - ) - ) - - self.features = nn.Sequential(*reversed(layers)) - - 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) -> Tensor: - x = self.features(x) - # x = torch.tanh(x) - return x - - def forward(self, x: Tensor) -> Tensor: - return self._forward_impl(x) - - def _convnext( block_setting: List[CNBlockConfig], stochastic_depth_prob: float, @@ -397,7 +232,7 @@ def _convnext( ) -> ConvNeXtBase: if reversed: - model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs) + model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs) else: model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)